// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Banded LU direct solver.
//!
//! A finite-element tangent on a structured mesh is banded: with the
//! flag's node numbering (short direction innermost) the 560-DOF Newton
//! tangent has a half-bandwidth of ~26, and the dense `LuDirect`
//! factorization — measured at 98% of the structural step, which is 79%
//! of a coupled FSI pass — does O(n³) work on entries that are known
//! zeros. This solver stores only the band (LAPACK `dgbtrf`-style
//! column-major band storage, `kl` extra rows for partial-pivoting
//! fill) and factorizes in O(n·kl·(kl+ku)).
//!
//! Partial pivoting rounds differently from `LuDirect`'s full pivoting,
//! so swapping solvers is a solver-path change: trajectories shift at
//! rounding level and must be re-verified against the pinned bands.

use super::{ConvergenceInfo, LinearSolver, SolverCapabilities, SolverOptions};
use crate::assembly::SparseMatrix;
use crate::error::{FeaResult, SolverError};
use nalgebra::DVector;
use std::time::Instant;

/// Banded LU with partial pivoting (row swaps confined to the band).
///
/// The band limits `kl`/`ku` are measured from the matrix handed to
/// `factorize` — a genuinely dense matrix degenerates to an unblocked
/// dense LU, so the solver is safe (if pointless) off the banded path.
#[derive(Debug, Default)]
pub struct BandedLu {
    /// Factor storage, column-major, `ldab = 2·kl + ku + 1` rows per
    /// column: row `kl + ku + i - j` of column `j` holds `A(i, j)`;
    /// the top `kl` rows are fill space for pivot swaps. Kept across
    /// calls so repeated same-size factorizations reuse the allocation.
    ab: Vec<f64>,
    ipiv: Vec<usize>,
    n: usize,
    kl: usize,
    ku: usize,
    factorized: bool,
}

impl BandedLu {
    /// Create a new banded LU solver.
    pub fn new() -> Self {
        Self::default()
    }

    /// The band limits `(kl, ku)` (sub- and super-diagonal counts) of
    /// the stored pattern, structural zeros included — the band is a
    /// property of the mesh topology, not of the current values.
    fn band_limits(matrix: &SparseMatrix) -> (usize, usize) {
        let n = matrix.nrows();
        let (row_ptr, col_idx) = matrix.structure();
        let mut kl = 0usize;
        let mut ku = 0usize;
        if row_ptr.len() == n + 1 {
            for row in 0..n {
                for &col in &col_idx[row_ptr[row]..row_ptr[row + 1]] {
                    if row > col {
                        kl = kl.max(row - col);
                    } else {
                        ku = ku.max(col - row);
                    }
                }
            }
        } else {
            // Not finalized to CSR — scan the dense image (cold path).
            let dense = matrix.to_dense();
            for row in 0..n {
                for col in 0..matrix.ncols() {
                    if dense[(row, col)] != 0.0 {
                        if row > col {
                            kl = kl.max(row - col);
                        } else {
                            ku = ku.max(col - row);
                        }
                    }
                }
            }
        }
        (kl, ku)
    }

    /// Factorize the matrix (band storage, partial pivoting).
    pub fn factorize(&mut self, matrix: &SparseMatrix) -> FeaResult<()> {
        let n = matrix.nrows();
        if n != matrix.ncols() {
            return Err(SolverError::FactorizationFailed {
                reason: format!("matrix is not square: {}x{}", n, matrix.ncols()),
            }
            .into());
        }
        let (kl, ku) = Self::band_limits(matrix);
        let ldab = 2 * kl + ku + 1;

        self.factorized = false;
        self.ab.clear();
        self.ab.resize(ldab * n, 0.0);
        self.ipiv.clear();
        self.ipiv.resize(n, 0);
        self.n = n;
        self.kl = kl;
        self.ku = ku;

        let ab = &mut self.ab;
        let band = |i: usize, j: usize| kl + ku + i - j + j * ldab;

        let (row_ptr, col_idx) = matrix.structure();
        if row_ptr.len() == n + 1 {
            let values = matrix.values();
            for row in 0..n {
                for idx in row_ptr[row]..row_ptr[row + 1] {
                    ab[band(row, col_idx[idx])] = values[idx];
                }
            }
        } else {
            let dense = matrix.to_dense();
            for row in 0..n {
                for col in 0..n {
                    let value = dense[(row, col)];
                    if value != 0.0 {
                        ab[band(row, col)] = value;
                    }
                }
            }
        }

        // Unblocked band factorization (LAPACK dgbtf2). U's bandwidth
        // grows to ku + kl from the row swaps; L's multipliers stay in
        // the kl rows under the diagonal of their column.
        for j in 0..n {
            let km = kl.min(n - 1 - j);
            let mut jp = 0usize;
            let mut pivot_abs = 0.0f64;
            for p in 0..=km {
                let a = ab[band(j + p, j)].abs();
                if a > pivot_abs {
                    pivot_abs = a;
                    jp = p;
                }
            }
            if pivot_abs == 0.0 {
                return Err(SolverError::FactorizationFailed {
                    reason: "banded LU factorization failed - matrix is singular".to_string(),
                }
                .into());
            }
            self.ipiv[j] = j + jp;
            let jw = (j + ku + kl).min(n - 1);
            if jp != 0 {
                for c in j..=jw {
                    ab.swap(band(j, c), band(j + jp, c));
                }
            }
            let pivot = ab[band(j, j)];
            for p in 1..=km {
                ab[band(j + p, j)] /= pivot;
            }
            for c in (j + 1)..=jw {
                let ujc = ab[band(j, c)];
                if ujc != 0.0 {
                    for p in 1..=km {
                        ab[band(j + p, c)] -= ab[band(j + p, j)] * ujc;
                    }
                }
            }
        }

        self.factorized = true;
        Ok(())
    }

    /// Solve with the current factors (forward with pivots, then back
    /// substitution through U's widened band).
    fn solve_factored(&self, rhs: &DVector<f64>) -> FeaResult<DVector<f64>> {
        if !self.factorized {
            return Err(SolverError::FactorizationRequired.into());
        }
        let (n, kl, ku) = (self.n, self.kl, self.ku);
        let ldab = 2 * kl + ku + 1;
        let ab = &self.ab;
        let band = |i: usize, j: usize| kl + ku + i - j + j * ldab;

        let mut x = rhs.clone();
        for j in 0..n {
            let jp = self.ipiv[j];
            if jp != j {
                x.swap_rows(j, jp);
            }
            let xj = x[j];
            if xj != 0.0 {
                for p in 1..=kl.min(n - 1 - j) {
                    x[j + p] -= ab[band(j + p, j)] * xj;
                }
            }
        }
        for j in (0..n).rev() {
            let xj = x[j] / ab[band(j, j)];
            x[j] = xj;
            if xj != 0.0 {
                for i in j.saturating_sub(ku + kl)..j {
                    x[i] -= ab[band(i, j)] * xj;
                }
            }
        }
        Ok(x)
    }
}

impl LinearSolver for BandedLu {
    fn solve(
        &mut self,
        matrix: &SparseMatrix,
        rhs: &DVector<f64>,
        _options: &SolverOptions,
    ) -> FeaResult<(DVector<f64>, ConvergenceInfo)> {
        let start_time = Instant::now();
        let mut info = ConvergenceInfo::new();

        if matrix.nrows() != rhs.len() {
            return Err(SolverError::DimensionMismatch {
                matrix_rows: matrix.nrows(),
                matrix_cols: matrix.ncols(),
                rhs_rows: rhs.len(),
                rhs_cols: 1,
            }
            .into());
        }

        // Always factorize the matrix we were handed (the Newton loop
        // changes values, never the size — see LuDirect's history).
        self.factorize(matrix)?;
        let solution = self.solve_factored(rhs)?;

        info.set_solve_time(start_time.elapsed());
        info.set_converged(1, 0.0, 0.0);
        info.set_memory_usage(self.ab.len() * 8);

        Ok((solution, info))
    }

    fn name(&self) -> &'static str {
        "Banded LU Direct"
    }

    fn capabilities(&self) -> SolverCapabilities {
        SolverCapabilities {
            symmetric: false,
            positive_definite: false,
            gpu_acceleration: false,
            multiple_rhs: true,
            iterative_refinement: true,
            memory_efficiency: 4,
            computational_efficiency: 5,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::solvers::LuDirect;

    /// A deterministic banded test matrix: diagonally dominant so the
    /// well-conditioned comparison against LuDirect is legitimate, with
    /// off-band entries exactly zero.
    fn banded_matrix(n: usize, kl: usize, ku: usize) -> SparseMatrix {
        let mut m = SparseMatrix::new(n, n);
        for i in 0..n {
            let mut off_sum = 0.0;
            for j in i.saturating_sub(kl)..=(i + ku).min(n - 1) {
                if i != j {
                    let v = ((7 * i + 13 * j + 3) as f64).sin();
                    m.add_entry(i, j, v).unwrap();
                    off_sum += v.abs();
                }
            }
            m.add_entry(i, i, off_sum + 1.0 + (i as f64 * 0.7).cos())
                .unwrap();
        }
        m.finalize().unwrap();
        m
    }

    fn manufactured_rhs(m: &SparseMatrix) -> (DVector<f64>, DVector<f64>) {
        let n = m.nrows();
        let x_exact = DVector::from_fn(n, |i, _| ((i as f64) * 0.31).sin() + 1.5);
        let b = m.multiply_vector(&x_exact).unwrap();
        (x_exact, b)
    }

    /// The core manufactured-solution check: recover a known x to
    /// near-machine precision, and agree with the dense LuDirect
    /// answer (different pivoting, same system).
    fn check_against_manufactured_and_dense(m: &SparseMatrix) {
        let (x_exact, b) = manufactured_rhs(m);
        let options = SolverOptions::default();

        let (x_banded, info) = BandedLu::new().solve(m, &b, &options).unwrap();
        assert!(info.converged);

        let rel_exact = (&x_banded - &x_exact).norm() / x_exact.norm();
        assert!(
            rel_exact < 1e-12,
            "banded solution off the manufactured x: rel err {rel_exact:.3e}"
        );

        let (x_dense, _) = LuDirect::new().solve(m, &b, &options).unwrap();
        let rel_dense = (&x_banded - &x_dense).norm() / x_dense.norm();
        assert!(
            rel_dense < 1e-12,
            "banded and dense LU disagree: rel err {rel_dense:.3e}"
        );
    }

    #[test]
    fn manufactured_symmetric_band() {
        check_against_manufactured_and_dense(&banded_matrix(60, 3, 3));
    }

    #[test]
    fn manufactured_asymmetric_band() {
        check_against_manufactured_and_dense(&banded_matrix(45, 5, 1));
        check_against_manufactured_and_dense(&banded_matrix(45, 1, 5));
    }

    #[test]
    fn manufactured_tridiagonal_and_diagonal() {
        check_against_manufactured_and_dense(&banded_matrix(30, 1, 1));
        check_against_manufactured_and_dense(&banded_matrix(12, 0, 0));
    }

    #[test]
    fn manufactured_full_bandwidth() {
        // kl = ku = n - 1: the band degenerates to dense storage and the
        // algorithm to an unblocked dense LU — must still be correct.
        check_against_manufactured_and_dense(&banded_matrix(10, 9, 9));
    }

    #[test]
    fn manufactured_single_dof() {
        check_against_manufactured_and_dense(&banded_matrix(1, 0, 0));
    }

    /// Zero diagonal forces a pivot swap on the very first column; an
    /// unpivoted band elimination fails here, a pivoted one must not.
    #[test]
    fn pivoting_zero_diagonal() {
        let mut m = SparseMatrix::new(2, 2);
        m.add_entry(0, 1, 1.0).unwrap();
        m.add_entry(1, 0, 1.0).unwrap();
        m.finalize().unwrap();
        let b = DVector::from_vec(vec![2.0, 3.0]);
        let (x, _) = BandedLu::new()
            .solve(&m, &b, &SolverOptions::default())
            .unwrap();
        assert!((x[0] - 3.0).abs() < 1e-14 && (x[1] - 2.0).abs() < 1e-14);
    }

    /// An indefinite symmetric system (a shifted stiffness — the shape
    /// of a Newton tangent near a turning point): no positive-definite
    /// shortcut may be assumed.
    #[test]
    fn indefinite_tangent_like_system() {
        let n = 40;
        let mut m = SparseMatrix::new(n, n);
        for i in 0..n {
            // 1-D stiffness [ -1, 2, -1 ] shifted by -3.2: eigenvalues
            // 2 - 2cos(kπ/(n+1)) - 3.2 straddle zero.
            m.add_entry(i, i, 2.0 - 3.2).unwrap();
            if i + 1 < n {
                m.add_entry(i, i + 1, -1.0).unwrap();
                m.add_entry(i + 1, i, -1.0).unwrap();
            }
        }
        m.finalize().unwrap();
        let (x_exact, b) = manufactured_rhs(&m);
        let (x, _) = BandedLu::new()
            .solve(&m, &b, &SolverOptions::default())
            .unwrap();
        let rel = (&x - &x_exact).norm() / x_exact.norm();
        assert!(rel < 1e-10, "indefinite solve rel err {rel:.3e}");
    }

    #[test]
    fn singular_matrix_is_detected() {
        let mut m = SparseMatrix::new(3, 3);
        // Row 2 is a copy of row 1 within the band.
        m.add_entry(0, 0, 2.0).unwrap();
        m.add_entry(0, 1, 1.0).unwrap();
        m.add_entry(1, 0, 4.0).unwrap();
        m.add_entry(1, 1, 3.0).unwrap();
        m.add_entry(2, 1, 3.0).unwrap();
        m.add_entry(2, 2, 0.0).unwrap();
        m.add_entry(1, 2, 0.0).unwrap();
        m.add_entry(2, 0, 4.0).unwrap();
        m.finalize().unwrap();
        let b = DVector::from_vec(vec![1.0, 1.0, 1.0]);
        assert!(
            BandedLu::new()
                .solve(&m, &b, &SolverOptions::default())
                .is_err()
        );
    }

    #[test]
    fn dimension_mismatch_is_rejected() {
        let m = banded_matrix(4, 1, 1);
        let b = DVector::from_vec(vec![1.0, 2.0]);
        assert!(
            BandedLu::new()
                .solve(&m, &b, &SolverOptions::default())
                .is_err()
        );
    }

    /// Two factorizations back to back through the same solver (the
    /// Newton pattern): the second must not see the first's factors.
    #[test]
    fn refactorizes_per_solve() {
        let options = SolverOptions::default();
        let mut solver = BandedLu::new();
        let m1 = banded_matrix(20, 2, 2);
        let (x1_exact, b1) = manufactured_rhs(&m1);
        let (x1, _) = solver.solve(&m1, &b1, &options).unwrap();
        assert!((&x1 - &x1_exact).norm() / x1_exact.norm() < 1e-12);

        let m2 = banded_matrix(20, 4, 3);
        let (x2_exact, b2) = manufactured_rhs(&m2);
        let (x2, _) = solver.solve(&m2, &b2, &options).unwrap();
        assert!((&x2 - &x2_exact).norm() / x2_exact.norm() < 1e-12);
    }
}