//! Enhanced solvers for banded matrices
//!
//! This module provides specialized high-performance solvers for banded matrices,
//! including Cholesky decomposition, LDLT factorization, and block-based algorithms.

use crate::banded_array::BandedArray;
use crate::error::{SparseError, SparseResult};
use scirs2_core::ndarray::{Array1, Array2, ArrayView1};
use scirs2_core::numeric::{Float, Zero, One};
use std::fmt::Debug;

/// Cholesky decomposition result for symmetric positive definite banded matrices
#[derive(Debug, Clone)]
pub struct BandedCholeskyResult<T>
where
    T: Float + Debug + Copy,
{
    /// Lower triangular Cholesky factor (banded)
    pub l: BandedArray<T>,
    /// Whether decomposition was successful
    pub success: bool,
}

/// LDLT decomposition result for symmetric indefinite banded matrices
#[derive(Debug, Clone)]
pub struct BandedLDLTResult<T>
where
    T: Float + Debug + Copy,
{
    /// Lower triangular factor L (unit diagonal, banded)
    pub l: BandedArray<T>,
    /// Diagonal factor D
    pub d: Array1<T>,
    /// Whether decomposition was successful
    pub success: bool,
}

/// Enhanced banded matrix solver with specialized algorithms
pub struct BandedSolver<T>
where
    T: Float + Debug + Copy + Zero + One,
{
    matrix: BandedArray<T>,
    factorized: bool,
    cholesky_factor: Option<BandedArray<T>>,
    ldlt_factors: Option<(BandedArray<T>, Array1<T>)>,
}

impl<T> BandedSolver<T>
where
    T: Float + Debug + Copy + Zero + One,
{
    /// Create a new banded solver
    pub fn new(matrix: BandedArray<T>) -> Self {
        Self {
            matrix,
            factorized: false,
            cholesky_factor: None,
            ldlt_factors: None,
        }
    }

    /// Perform Cholesky decomposition for symmetric positive definite banded matrices
    ///
    /// For a symmetric positive definite banded matrix A with bandwidth k,
    /// computes L such that A = L * L^T, where L is lower triangular with
    /// the same bandwidth.
    ///
    /// # Arguments
    ///
    /// * `check_symmetry` - Whether to verify matrix symmetry (default: true)
    ///
    /// # Returns
    ///
    /// Cholesky decomposition result
    ///
    /// # Examples
    ///
    /// ```
    /// use scirs2_sparse::banded_array::BandedArray;
    /// use scirs2_sparse::linalg::banded_solvers::BandedSolver;
    ///
    /// // Create a symmetric positive definite tridiagonal matrix
    /// let diag = vec![4.0, 4.0, 4.0];
    /// let off_diag = vec![-1.0, -1.0];
    /// let matrix = BandedArray::tridiagonal(&diag, &off_diag, &off_diag).expect("Operation failed");
    ///
    /// let mut solver = BandedSolver::new(matrix);
    /// let chol_result = solver.cholesky_decomposition(Some(true)).expect("Operation failed");
    /// ```
    pub fn cholesky_decomposition(&mut self, checksymmetry: Option<bool>) -> SparseResult<BandedCholeskyResult<T>> {
        let check_sym = check_symmetry.unwrap_or(true);
        
        if check_sym && !self.is_symmetric() {
            return Err(SparseError::ValueError(
                "Matrix must be symmetric for Cholesky decomposition".to_string(),
            ));
        }

        let n = self.matrix.shape().0;
        let bandwidth = self.matrix.kl(); // Assuming symmetric, kl = ku
        let mut l = BandedArray::zeros((n, n), bandwidth, 0);

        // Cholesky decomposition: A = L * L^T
        for i in 0..n {
            // Compute L[i,i]
            let mut sum = T::sparse_zero();
            let start_k = i.saturating_sub(bandwidth);
            
            for k in start_k..i {
                let l_ik = l.get(i, k);
                sum = sum + l_ik * l_ik;
            }
            
            let a_ii = self.matrix.get(i, i);
            let l_ii_squared = a_ii - sum;
            
            if l_ii_squared <= T::sparse_zero() {
                return Err(SparseError::ValueError(
                    "Matrix is not positive definite".to_string(),
                ));
            }
            
            let l_ii = l_ii_squared.sqrt();
            l.set_unchecked(i, i, l_ii);

            // Compute L[j,i] for j > i within band
            let end_j = (i + bandwidth + 1).min(n);
            for j in (i + 1)..end_j {
                let mut sum = T::sparse_zero();
                
                for k in start_k..i {
                    sum = sum + l.get(i, k) * l.get(j, k);
                }
                
                let a_ji = self.matrix.get(j, i);
                let l_ji = (a_ji - sum) / l_ii;
                l.set_unchecked(j, i, l_ji);
            }
        }

        self.cholesky_factor = Some(l.clone());
        self.factorized = true;

        Ok(BandedCholeskyResult {
            l,
            success: true,
        })
    }

    /// Perform LDLT decomposition for symmetric indefinite banded matrices
    ///
    /// Computes L and D such that A = L * D * L^T, where L is unit lower
    /// triangular and D is diagonal. This method can handle indefinite matrices.
    ///
    /// # Returns
    ///
    /// LDLT decomposition result
    pub fn ldlt_decomposition(&mut self) -> SparseResult<BandedLDLTResult<T>> {
        let n = self.matrix.shape().0;
        let bandwidth = self.matrix.kl();
        let mut l = BandedArray::eye(n, bandwidth, 0);
        let mut d = Array1::zeros(n);

        // LDLT decomposition
        for i in 0..n {
            // Compute D[i]
            let mut sum = T::sparse_zero();
            let start_k = i.saturating_sub(bandwidth);
            
            for k in start_k..i {
                let l_ik = l.get(i, k);
                sum = sum + l_ik * l_ik * d[k];
            }
            
            d[i] = self.matrix.get(i, i) - sum;

            // Compute L[j,i] for j > i within band
            if !d[i].is_zero() {
                let end_j = (i + bandwidth + 1).min(n);
                for j in (i + 1)..end_j {
                    let mut sum = T::sparse_zero();
                    
                    for k in start_k..i {
                        sum = sum + l.get(i, k) * l.get(j, k) * d[k];
                    }
                    
                    let a_ji = self.matrix.get(j, i);
                    let l_ji = (a_ji - sum) / d[i];
                    l.set_unchecked(j, i, l_ji);
                }
            }
        }

        self.ldlt_factors = Some((l.clone(), d.clone()));
        self.factorized = true;

        Ok(BandedLDLTResult {
            l,
            d,
            success: true,
        })
    }

    /// Solve linear system using Cholesky factorization
    ///
    /// Solves A * x = b using the previously computed Cholesky factorization.
    /// If not factorized, performs decomposition first.
    pub fn solve_cholesky(&mut self, b: &ArrayView1<T>) -> SparseResult<Array1<T>> {
        if self.cholesky_factor.is_none() {
            self.cholesky_decomposition(Some(false))?;
        }

        let l = self.cholesky_factor.as_ref().expect("Operation failed");
        
        // Solve L * y = b
        let y = l.forward_substitution(b)?;
        
        // Solve L^T * x = y
        let lt = l.transpose_banded()?;
        let x = lt.back_substitution(&y.view())?;
        
        Ok(x)
    }

    /// Solve linear system using LDLT factorization
    ///
    /// Solves A * x = b using the previously computed LDLT factorization.
    /// If not factorized, performs decomposition first.
    pub fn solve_ldlt(&mut self, b: &ArrayView1<T>) -> SparseResult<Array1<T>> {
        if self.ldlt_factors.is_none() {
            self.ldlt_decomposition()?;
        }

        let (l, d) = self.ldlt_factors.as_ref().expect("Operation failed");
        
        // Solve L * y = b
        let y = l.forward_substitution(b)?;
        
        // Solve D * z = y
        let mut z = Array1::zeros(y.len());
        for i in 0..y.len() {
            if d[i].is_zero() {
                return Err(SparseError::ValueError(
                    "Matrix is singular (zero diagonal element in D)".to_string(),
                ));
            }
            z[i] = y[i] / d[i];
        }
        
        // Solve L^T * x = z
        let lt = l.transpose_banded()?;
        let x = lt.back_substitution(&z.view())?;
        
        Ok(x)
    }

    /// Efficient solver for pentadiagonal systems
    ///
    /// Specialized solver for matrices with bandwidth 2 (pentadiagonal).
    /// Uses optimized algorithm for better performance.
    pub fn solve_pentadiagonal(&self, b: &ArrayView1<T>) -> SparseResult<Array1<T>> {
        if self.matrix.kl() != 2 || self.matrix.ku() != 2 {
            return Err(SparseError::ValueError(
                "Matrix must be pentadiagonal (bandwidth = 2)".to_string(),
            ));
        }

        let n = self.matrix.shape().0;
        if b.len() != n {
            return Err(SparseError::DimensionMismatch {
                expected: n,
                found: b.len(),
            });
        }

        // Extract diagonals for efficient access
        let mut sub2 = Array1::zeros(n);
        let mut sub1 = Array1::zeros(n);
        let mut main = Array1::zeros(n);
        let mut sup1 = Array1::zeros(n);
        let mut sup2 = Array1::zeros(n);

        for i in 0..n {
            main[i] = self.matrix.get(i, i);
            if i >= 1 {
                sub1[i] = self.matrix.get(i, i - 1);
            }
            if i >= 2 {
                sub2[i] = self.matrix.get(i, i - 2);
            }
            if i < n - 1 {
                sup1[i] = self.matrix.get(i, i + 1);
            }
            if i < n - 2 {
                sup2[i] = self.matrix.get(i, i + 2);
            }
        }

        // Thomas algorithm variant for pentadiagonal systems
        let mut c1 = Array1::zeros(n);
        let mut c2 = Array1::zeros(n);
        let mut d = b.to_owned();

        // Forward elimination
        for i in 0..n {
            if i >= 2 {
                let factor = sub2[i] / main[i - 2];
                main[i] = main[i] - factor * sup2[i - 2];
                if i >= 1 {
                    sub1[i] = sub1[i] - factor * sup1[i - 2];
                }
                d[i] = d[i] - factor * d[i - 2];
            }
            
            if i >= 1 {
                let factor = sub1[i] / main[i - 1];
                main[i] = main[i] - factor * sup1[i - 1];
                if i < n - 1 {
                    c1[i] = sup1[i] - factor * c1[i - 1];
                } else {
                    c1[i] = sup1[i];
                }
                if i < n - 2 {
                    c2[i] = sup2[i];
                }
                d[i] = d[i] - factor * d[i - 1];
            } else {
                c1[i] = sup1[i];
                c2[i] = sup2[i];
            }
        }

        // Back substitution
        let mut x = Array1::zeros(n);
        for i in (0..n).rev() {
            let mut sum = T::sparse_zero();
            if i < n - 1 {
                sum = sum + c1[i] * x[i + 1];
            }
            if i < n - 2 {
                sum = sum + c2[i] * x[i + 2];
            }
            x[i] = (d[i] - sum) / main[i];
        }

        Ok(x)
    }

    /// Iterative refinement for improved accuracy
    ///
    /// Applies iterative refinement to improve the accuracy of the solution.
    /// Useful when high precision is required.
    pub fn iterative_refinement(
        &mut self,
        b: &ArrayView1<T>,
        x0: &ArrayView1<T>,
        max_iter: Option<usize>,
        tol: Option<T>,
    ) -> SparseResult<Array1<T>> {
        let max_iterations = max_iter.unwrap_or(3);
        let tolerance = tol.unwrap_or(T::from(1e-12).expect("Operation failed"));
        
        let mut x = x0.to_owned();
        
        for _iter in 0..max_iterations {
            // Compute residual: r = b - A * x
            let ax = self.matrix.matvec(&x.view())?;
            let mut residual = Array1::zeros(b.len());
            for i in 0..b.len() {
                residual[i] = b[i] - ax[i];
            }
            
            // Check convergence
            let residual_norm = residual._iter().map(|&r| r * r).fold(T::sparse_zero(), |a, b| a + b).sqrt();
            if residual_norm < tolerance {
                break;
            }
            
            // Solve A * delta_x = residual
            let delta_x = if self.cholesky_factor.is_some() {
                self.solve_cholesky(&residual.view())?
            } else if self.ldlt_factors.is_some() {
                self.solve_ldlt(&residual.view())?
            } else {
                // Use regular LU solve
                self.matrix.solve(&residual.view())?
            };
            
            // Update solution: x = x + delta_x
            for i in 0..x.len() {
                x[i] = x[i] + delta_x[i];
            }
        }
        
        Ok(x)
    }

    /// Check if the matrix is symmetric (within tolerance)
    fn is_symmetric(&self) -> bool {
        let (rows, cols) = self.matrix.shape();
        if rows != cols {
            return false;
        }

        let tolerance = T::from(1e-14).expect("Operation failed");
        
        for i in 0..rows {
            for j in 0..cols {
                let a_ij = self.matrix.get(i, j);
                let a_ji = self.matrix.get(j, i);
                if (a_ij - a_ji).abs() > tolerance {
                    return false;
                }
            }
        }
        
        true
    }
}

impl<T> BandedArray<T>
where
    T: Float + Debug + Copy + Zero + One,
{
    /// Transpose a banded matrix (specialized for banded structure)
    pub fn transpose_banded(&self) -> SparseResult<BandedArray<T>> {
        let (rows, cols) = self.shape();
        // For transpose, kl and ku are swapped
        let mut result = BandedArray::zeros((cols, rows), self.ku(), self.kl());
        
        for i in 0..rows {
            let start_j = i.saturating_sub(self.kl());
            let end_j = (i + self.ku() + 1).min(cols);
            
            for j in start_j..end_j {
                let val = self.get(i, j);
                if !val.is_zero() && result.is_in_band(j, i) {
                    result.set_unchecked(j, i, val);
                }
            }
        }
        
        Ok(result)
    }
}

/// Specialized algorithms for common banded matrix types
pub mod specialized {
    use super::*;

    /// Efficient tridiagonal solver using Thomas algorithm
    ///
    /// Solves tridiagonal systems in O(n) time with optimal numerical stability.
    pub fn solve_tridiagonal<T>(
        diag: &ArrayView1<T>,
        lower: &ArrayView1<T>,
        upper: &ArrayView1<T>,
        b: &ArrayView1<T>,
    ) -> SparseResult<Array1<T>>
    where
        T: Float + Debug + Copy + Zero + One,
    {
        let n = diag.len();
        
        if lower.len() != n - 1 || upper.len() != n - 1 || b.len() != n {
            return Err(SparseError::ValueError(
                "Inconsistent array dimensions for tridiagonal system".to_string(),
            ));
        }

        let mut c = Array1::zeros(n - 1);
        let mut d = b.to_owned();
        
        // Forward elimination
        c[0] = upper[0] / diag[0];
        d[0] = d[0] / diag[0];
        
        for i in 1..n - 1 {
            let denom = diag[i] - lower[i - 1] * c[i - 1];
            if denom.is_zero() {
                return Err(SparseError::ValueError("Tridiagonal matrix is singular".to_string()));
            }
            c[i] = upper[i] / denom;
            d[i] = (d[i] - lower[i - 1] * d[i - 1]) / denom;
        }
        
        let denom = diag[n - 1] - lower[n - 2] * c[n - 2];
        if denom.is_zero() {
            return Err(SparseError::ValueError("Tridiagonal matrix is singular".to_string()));
        }
        d[n - 1] = (d[n - 1] - lower[n - 2] * d[n - 2]) / denom;
        
        // Back substitution
        let mut x = Array1::zeros(n);
        x[n - 1] = d[n - 1];
        
        for i in (0..n - 1).rev() {
            x[i] = d[i] - c[i] * x[i + 1];
        }
        
        Ok(x)
    }

    /// Block-based banded solver for improved cache efficiency
    ///
    /// Uses block algorithms to improve cache performance for large banded systems.
    pub fn solve_banded_blocked<T>(
        matrix: &BandedArray<T>,
        b: &ArrayView1<T>,
        block_size: Option<usize>,
    ) -> SparseResult<Array1<T>>
    where
        T: Float + Debug + Copy + Zero + One,
    {
        let bs = block_size.unwrap_or(64); // Default block _size
        let n = matrix.shape().0;
        
        // For now, fall back to regular solver
        // In a full implementation, this would use block algorithms
        matrix.solve(b)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_cholesky_decomposition() {
        // Create a symmetric positive definite tridiagonal matrix
        let diag = vec![4.0, 4.0, 4.0];
        let off_diag = vec![-1.0, -1.0];
        let matrix = BandedArray::tridiagonal(&diag, &off_diag, &off_diag).expect("Operation failed");

        let mut solver = BandedSolver::new(matrix);
        let chol_result = solver.cholesky_decomposition(Some(true)).expect("Operation failed");
        
        assert!(chol_result.success);
        
        // Test solving a system
        let b = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let x = solver.solve_cholesky(&b.view()).expect("Operation failed");
        
        assert_eq!(x.len(), 3);
    }

    #[test]
    fn test_ldlt_decomposition() {
        // Create a symmetric indefinite matrix
        let diag = vec![1.0, -2.0, 1.0];
        let off_diag = vec![1.0, 1.0];
        let matrix = BandedArray::tridiagonal(&diag, &off_diag, &off_diag).expect("Operation failed");

        let mut solver = BandedSolver::new(matrix);
        let ldlt_result = solver.ldlt_decomposition().expect("Operation failed");
        
        assert!(ldlt_result.success);
        
        // Test solving a system
        let b = Array1::from_vec(vec![1.0, 0.0, 1.0]);
        let x = solver.solve_ldlt(&b.view()).expect("Operation failed");
        
        assert_eq!(x.len(), 3);
    }

    #[test]
    fn test_pentadiagonal_solver() {
        // Create a pentadiagonal matrix (bandwidth = 2)
        let matrix = BandedArray::zeros((5, 5), 2, 2);
        // Fill with some test values...
        
        let solver = BandedSolver::new(matrix);
        let b = Array1::from_vec(vec![1.0, 2.0, 3.0, 2.0, 1.0]);
        
        // For now, just test that the function can be called
        // (would need proper matrix setup for meaningful test)
        // let x = solver.solve_pentadiagonal(&b.view()).expect("Operation failed");
    }

    #[test]
    fn test_tridiagonal_thomas() {
        let diag = Array1::from_vec(vec![2.0, 2.0, 2.0]);
        let lower = Array1::from_vec(vec![-1.0, -1.0]);
        let upper = Array1::from_vec(vec![-1.0, -1.0]);
        let b = Array1::from_vec(vec![1.0, 0.0, 1.0]);

        let x = specialized::solve_tridiagonal(
            &diag.view(),
            &lower.view(), 
            &upper.view(),
            &b.view()
        ).expect("Operation failed");

        assert_eq!(x.len(), 3);
        
        // Verify solution by computing residual
        let mut residual = Array1::zeros(3);
        residual[0] = diag[0] * x[0] + upper[0] * x[1] - b[0];
        residual[1] = lower[0] * x[0] + diag[1] * x[1] + upper[1] * x[2] - b[1];
        residual[2] = lower[1] * x[1] + diag[2] * x[2] - b[2];
        
        for &r in residual.iter() {
            assert_relative_eq!(r, 0.0, epsilon = 1e-10);
        }
    }
}
