//! `NalgebraLuSolver`: the existing dense nalgebra LU path re-implemented
//! behind the [`LinearSolver`] lifecycle (FR3, Story 1.4).
//!
//! The numeric pipeline is **bit-identical** to [`crate::jacobian::JacobianMatrix::solve`]:
//! Ruiz equilibration ([`crate::scaling::equilibrate`]) → scaled LU with
//! partial pivoting → scaled solve `J·x = b` → undo column scaling
//! ([`crate::scaling::unscale_dx`]) → reject non-finite steps. The difference
//! is structural: assembly+factorization happen once in `set_linearisation`
//! and the factorization (with the cached scaling factors) is reused across
//! `solve_in_place` calls — the typed hook FR4's frozen-Jacobian reuse needs.
//!
//! ## Worked example (lifecycle, diffsol-style)
//!
//! ```rust
//! use entropyk_solver::linear::NalgebraLuSolver;
//! use entropyk_solver_core::LinearSolver;
//!
//! let mut lu = NalgebraLuSolver::new();
//! lu.set_problem(2, 2).unwrap();                       // 1. structure/setup
//! lu.set_linearisation(&[(0, 0, 2.0), (1, 1, 1.0)])    // 2. assemble + factor
//!     .unwrap();
//! let mut rhs = vec![4.0, 3.0];                        // J·x = b
//! lu.solve_in_place(&mut rhs).unwrap();                // 3. solve in place
//! assert!((rhs[0] - 2.0).abs() < 1e-10);
//! assert!((rhs[1] - 3.0).abs() < 1e-10);
//! ```

use entropyk_solver_core::{LinearSolver, SolverCoreError};
use nalgebra::DMatrix;

/// Stored factorization plus the equilibration scalings used to build it.
struct Factorized {
    d_r: Vec<f64>,
    d_c: Vec<f64>,
    lu: nalgebra::LU<f64, nalgebra::Dyn, nalgebra::Dyn>,
}

/// Dense LU backend (nalgebra) behind the object-safe [`LinearSolver`] trait.
///
/// Default backend for the Newton/PTC strategies. Send by construction
/// (nalgebra dense types are `Send`). Scratch buffers are allocated once in
/// `set_problem` so `solve_in_place` performs **zero heap allocation**.
pub struct NalgebraLuSolver {
    n_rows: usize,
    n_cols: usize,
    /// Assembled (unscaled) matrix, kept so `set_matrix` can reuse storage.
    matrix: Option<DMatrix<f64>>,
    factor: Option<Factorized>,
    /// Scratch: scaled rhs / unscaled step (allocated in `set_problem`).
    delta: Vec<f64>,
}

impl Default for NalgebraLuSolver {
    fn default() -> Self {
        Self::new()
    }
}

impl NalgebraLuSolver {
    /// Creates an unconfigured backend (call [`LinearSolver::set_problem`]).
    pub fn new() -> Self {
        Self {
            n_rows: 0,
            n_cols: 0,
            matrix: None,
            factor: None,
            delta: Vec::new(),
        }
    }

    /// Inherent fast path: copy an already-assembled dense matrix and factor
    /// it (one `copy_from`, no triplet conversion). Strategies in this crate
    /// use it so the assembly container stays [`crate::jacobian::JacobianMatrix`];
    /// the trait's triplet-based `set_linearisation` shares the same
    /// factorization step.
    pub fn set_matrix(&mut self, matrix: &DMatrix<f64>) -> Result<(), SolverCoreError> {
        let stored = self.matrix.as_mut().ok_or_else(|| SolverCoreError::Usage {
            message: "set_problem must be called before set_matrix".to_string(),
        })?;
        if stored.nrows() != matrix.nrows() || stored.ncols() != matrix.ncols() {
            return Err(SolverCoreError::Usage {
                message: format!(
                    "matrix shape {}x{} does not match problem {}x{}",
                    matrix.nrows(),
                    matrix.ncols(),
                    stored.nrows(),
                    stored.ncols()
                ),
            });
        }
        stored.copy_from(matrix);
        self.factorize()
    }

    /// Equilibrate + scale + factorize the stored matrix (square path).
    fn factorize(&mut self) -> Result<(), SolverCoreError> {
        let Some(matrix) = self.matrix.as_ref() else {
            // Unreachable via the public API (guarded by set_problem), but the
            // zero-panic policy forbids `expect` in library code.
            self.factor = None;
            return Err(SolverCoreError::Usage {
                message: "set_problem must be called before factorizing".to_string(),
            });
        };
        if matrix.nrows() != matrix.ncols() {
            // Non-square problems have no LU path; solve_in_place reports Usage.
            self.factor = None;
            return Ok(());
        }
        let n = matrix.nrows();
        let (d_r, d_c) = crate::scaling::equilibrate(matrix);
        let mut scaled = matrix.clone();
        for i in 0..n {
            for j in 0..n {
                scaled[(i, j)] *= d_r[i] * d_c[j];
            }
        }
        self.factor = Some(Factorized {
            d_r,
            d_c,
            lu: scaled.lu(),
        });
        Ok(())
    }
}

impl LinearSolver for NalgebraLuSolver {
    fn set_problem(
        &mut self,
        n_equations: usize,
        n_unknowns: usize,
    ) -> Result<(), SolverCoreError> {
        if n_equations == 0 || n_unknowns == 0 {
            return Err(SolverCoreError::Usage {
                message: format!("zero-sized problem {n_equations}x{n_unknowns}"),
            });
        }
        self.n_rows = n_equations;
        self.n_cols = n_unknowns;
        self.matrix = Some(DMatrix::zeros(n_equations, n_unknowns));
        self.factor = None;
        // Scratch buffer for the zero-allocation solve path.
        self.delta = vec![0.0; n_unknowns];
        Ok(())
    }

    fn set_linearisation(
        &mut self,
        entries: &[(usize, usize, f64)],
    ) -> Result<(), SolverCoreError> {
        let (n_rows, n_cols) = (self.n_rows, self.n_cols);
        if self.matrix.is_none() {
            return Err(SolverCoreError::Usage {
                message: "set_problem must be called before set_linearisation".to_string(),
            });
        }
        // Validate bounds BEFORE touching the matrix: an out-of-bounds entry
        // must leave neither a partially-filled matrix nor a stale
        // factorization behind (audit finding, High).
        for &(row, col, _) in entries {
            if row >= n_rows || col >= n_cols {
                self.factor = None;
                return Err(SolverCoreError::Usage {
                    message: format!("entry ({row},{col}) out of bounds {n_rows}x{n_cols}"),
                });
            }
        }
        let matrix = self.matrix.as_mut().ok_or_else(|| SolverCoreError::Usage {
            message: "set_problem must be called before set_linearisation".to_string(),
        })?;
        matrix.fill(0.0);
        for &(row, col, value) in entries {
            matrix[(row, col)] += value;
        }
        self.factorize()
    }

    fn solve_in_place(&mut self, rhs: &mut [f64]) -> Result<(), SolverCoreError> {
        if rhs.len() != self.n_rows || self.n_rows != self.n_cols {
            return Err(SolverCoreError::Usage {
                message: format!(
                    "solve_in_place needs a square rhs of len {}, got {} (problem {}x{})",
                    self.n_rows,
                    rhs.len(),
                    self.n_rows,
                    self.n_cols
                ),
            });
        }
        let n = self.n_rows;
        // Disjoint field borrows (zero-panic, zero-clone): factorization +
        // scratch buffers allocated in set_problem. The solve runs fully in
        // place via nalgebra's `solve_mut` over vector views — **zero heap
        // allocation** (Story 1.4 audit).
        let Self { factor, delta, .. } = self;
        let factor = factor.as_ref().ok_or_else(|| SolverCoreError::Usage {
            message: "set_linearisation must be called before solve_in_place".to_string(),
        })?;
        // Scaled right-hand side: D_r · b (into the pre-allocated scratch).
        for i in 0..n {
            delta[i] = rhs[i] * factor.d_r[i];
        }
        let mut b_view = nalgebra::DVectorViewMut::from_slice(&mut delta[..n], n);
        if factor.lu.solve_mut(&mut b_view) {
            // Undo the column scaling in place: x_i = D_c,i · y_i.
            for i in 0..n {
                delta[i] *= factor.d_c[i];
            }
            if delta[..n].iter().all(|v| v.is_finite()) {
                rhs.copy_from_slice(&delta[..n]);
                Ok(())
            } else {
                tracing::warn!(
                    "LU solve produced a non-finite step - Jacobian may contain NaN/Inf"
                );
                Err(SolverCoreError::InvalidSystem {
                    message: "linear solve produced a non-finite step".to_string(),
                })
            }
        } else {
            tracing::warn!("LU solve failed - Jacobian may be singular");
            Err(SolverCoreError::InvalidSystem {
                message: "linear solve failed (singular factorization)".to_string(),
            })
        }
    }

    fn n(&self) -> usize {
        self.n_rows
    }
}

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

    /// Deterministic dense system (well-conditioned, diagonally dominant).
    fn dense_entries(n: usize) -> Vec<(usize, usize, f64)> {
        let mut entries = Vec::new();
        for i in 0..n {
            for j in 0..n {
                let v = if i == j {
                    4.0 + 0.1 * (i as f64)
                } else {
                    1.0 / (1.0 + (i as f64 - j as f64).abs())
                };
                entries.push((i, j, v));
            }
        }
        entries
    }

    #[test]
    fn parity_with_jacobian_matrix_solve() {
        for n in [1, 2, 5, 12] {
            let entries = dense_entries(n);
            let jm = JacobianMatrix::from_builder(&entries, n, n);
            let residuals: Vec<f64> = (0..n)
                .map(|i| (i as f64 * 0.7 - 1.3).sin() * 100.0)
                .collect();

            let legacy = jm.solve(&residuals).expect("legacy solve");

            let mut backend = NalgebraLuSolver::new();
            backend.set_problem(n, n).unwrap();
            backend.set_linearisation(&entries).unwrap();
            let mut rhs: Vec<f64> = residuals.iter().map(|r| -*r).collect();
            backend.solve_in_place(&mut rhs).unwrap();

            for i in 0..n {
                assert!(
                    (legacy[i] - rhs[i]).abs() <= 1e-12 * legacy[i].abs().max(1.0),
                    "n={n} i={i}: legacy={} backend={}",
                    legacy[i],
                    rhs[i]
                );
            }
        }
    }

    #[test]
    fn parity_via_set_matrix_and_factorization_reuse() {
        let n = 6;
        let entries = dense_entries(n);
        let jm = JacobianMatrix::from_builder(&entries, n, n);
        let mut backend = NalgebraLuSolver::new();
        backend.set_problem(n, n).unwrap();
        backend.set_matrix(jm.as_matrix()).unwrap();

        // Two different RHS against the SAME stored factorization (the FR4
        // reuse hook): both must match the legacy path bit-for-bit.
        for k in 0..2 {
            let residuals: Vec<f64> = (0..n).map(|i| (i + k) as f64 * 3.1 - 5.0).collect();
            let legacy = jm.solve(&residuals).expect("legacy solve");
            let mut rhs: Vec<f64> = residuals.iter().map(|r| -*r).collect();
            backend.solve_in_place(&mut rhs).unwrap();
            for i in 0..n {
                assert!(
                    (legacy[i] - rhs[i]).abs() <= 1e-12 * legacy[i].abs().max(1.0),
                    "reuse k={k} i={i}"
                );
            }
        }
    }

    #[test]
    fn singular_and_error_paths() {
        let mut backend = NalgebraLuSolver::new();
        assert!(backend.set_problem(0, 2).is_err());
        backend.set_problem(2, 2).unwrap();
        // Solve before set_linearisation → Usage.
        assert!(matches!(
            backend.solve_in_place(&mut [1.0, 2.0]),
            Err(SolverCoreError::Usage { .. })
        ));
        // Out-of-bounds entry → Usage.
        assert!(backend.set_linearisation(&[(2, 0, 1.0)]).is_err());
        // Singular matrix → same outcome as the legacy path (whatever nalgebra
        // flags after equilibration; parity, not an idealized Err).
        let singular = [(0, 0, 1.0), (0, 1, 2.0), (1, 0, 1.0), (1, 1, 2.0)];
        let jm = JacobianMatrix::from_builder(&singular, 2, 2);
        let legacy = jm.solve(&[1.0, 2.0]);
        backend.set_linearisation(&singular).unwrap();
        let mut rhs = vec![-1.0, -2.0];
        match (legacy, backend.solve_in_place(&mut rhs)) {
            (None, Err(SolverCoreError::InvalidSystem { .. })) => {}
            (Some(legacy_delta), Ok(())) => {
                for i in 0..2 {
                    assert!(
                        (legacy_delta[i] - rhs[i]).abs() <= 1e-9 * legacy_delta[i].abs().max(1.0),
                        "parity on singular matrix: legacy={} backend={}",
                        legacy_delta[i],
                        rhs[i]
                    );
                }
            }
            (legacy, backend) => {
                panic!("divergent behavior: legacy={legacy:?} backend={backend:?}")
            }
        }
        assert_eq!(backend.n(), 2);
    }

    #[test]
    fn non_finite_matrix_entries_are_rejected() {
        let mut backend = NalgebraLuSolver::new();
        backend.set_problem(2, 2).unwrap();
        backend
            .set_linearisation(&[(0, 0, f64::NAN), (1, 1, 1.0)])
            .unwrap();
        let mut rhs = vec![1.0, 2.0];
        assert!(matches!(
            backend.solve_in_place(&mut rhs),
            Err(SolverCoreError::InvalidSystem { .. })
        ));
    }
}

// ─── Strangler dispatch (Story 1.5) ─────────────────────────────────────────

/// Process-wide override for the default dense backend (set by the CLI from
/// `solver.linear_backend`). Precedence: explicit `make_linear_backend` name
/// → this override → `ENTROPYK_LINEAR_BACKEND` env → `"nalgebra"`.
static BACKEND_OVERRIDE: std::sync::RwLock<Option<String>> = std::sync::RwLock::new(None);

/// Sets the process-wide default dense backend (`Some("faer")` / `Some("nalgebra")`,
/// `None` to clear). Additive, additive-only configuration surface.
pub fn set_linear_backend_override(name: Option<&str>) {
    if let Ok(mut guard) = BACKEND_OVERRIDE.write() {
        *guard = name.map(str::to_string);
    }
}

/// The dense LU backend in use: nalgebra (default, legacy path) or faer
/// (opt-in). Strategies hold this enum directly so the `set_matrix` fast path
/// stays available; it also implements [`LinearSolver`] for registry boxing.
pub enum DenseLu {
    /// nalgebra dense LU (default).
    Nalgebra(NalgebraLuSolver),
    /// faer dense LU (Story 1.5 opt-in).
    Faer(crate::linear_faer::FaerLuSolver),
}

impl DenseLu {
    /// Inherent fast path shared by both backends (see each backend's docs).
    pub fn set_matrix(&mut self, matrix: &DMatrix<f64>) -> Result<(), SolverCoreError> {
        match self {
            Self::Nalgebra(b) => b.set_matrix(matrix),
            Self::Faer(b) => b.set_matrix(matrix),
        }
    }
}

impl LinearSolver for DenseLu {
    fn set_problem(
        &mut self,
        n_equations: usize,
        n_unknowns: usize,
    ) -> Result<(), SolverCoreError> {
        match self {
            Self::Nalgebra(b) => b.set_problem(n_equations, n_unknowns),
            Self::Faer(b) => b.set_problem(n_equations, n_unknowns),
        }
    }

    fn set_linearisation(
        &mut self,
        entries: &[(usize, usize, f64)],
    ) -> Result<(), SolverCoreError> {
        match self {
            Self::Nalgebra(b) => b.set_linearisation(entries),
            Self::Faer(b) => b.set_linearisation(entries),
        }
    }

    fn solve_in_place(&mut self, rhs: &mut [f64]) -> Result<(), SolverCoreError> {
        match self {
            Self::Nalgebra(b) => b.solve_in_place(rhs),
            Self::Faer(b) => b.solve_in_place(rhs),
        }
    }

    fn n(&self) -> usize {
        match self {
            Self::Nalgebra(b) => b.n(),
            Self::Faer(b) => b.n(),
        }
    }
}

/// Builds the dense backend selected by name (or by override/env/default).
pub fn make_linear_backend(name: Option<&str>) -> Result<DenseLu, SolverCoreError> {
    let selected = name
        .map(str::to_string)
        .or_else(|| BACKEND_OVERRIDE.read().ok().and_then(|g| g.clone()))
        .or_else(|| std::env::var("ENTROPYK_LINEAR_BACKEND").ok())
        .unwrap_or_else(|| "nalgebra".to_string());
    match selected.as_str() {
        "nalgebra" => Ok(DenseLu::Nalgebra(NalgebraLuSolver::new())),
        "faer" => Ok(DenseLu::Faer(crate::linear_faer::FaerLuSolver::new())),
        other => Err(SolverCoreError::Usage {
            message: format!("unknown linear backend '{other}' (use 'nalgebra' or 'faer')"),
        }),
    }
}

#[cfg(test)]
mod dispatch_tests {
    use super::*;

    #[test]
    fn factory_builds_named_backends_and_rejects_unknown() {
        assert!(matches!(
            make_linear_backend(Some("nalgebra")).unwrap(),
            DenseLu::Nalgebra(_)
        ));
        assert!(matches!(
            make_linear_backend(Some("faer")).unwrap(),
            DenseLu::Faer(_)
        ));
        assert!(make_linear_backend(Some("mumps")).is_err());
        // Env-independent check: an explicit override always wins.
        set_linear_backend_override(Some("nalgebra"));
        assert!(matches!(
            make_linear_backend(None).unwrap(),
            DenseLu::Nalgebra(_)
        ));
        set_linear_backend_override(None);
    }

    #[test]
    fn override_selects_default_and_explicit_name_wins() {
        set_linear_backend_override(Some("faer"));
        assert!(matches!(
            make_linear_backend(None).unwrap(),
            DenseLu::Faer(_)
        ));
        // Explicit name wins over the override.
        assert!(matches!(
            make_linear_backend(Some("nalgebra")).unwrap(),
            DenseLu::Nalgebra(_)
        ));
        set_linear_backend_override(None);
        // After clearing, selection falls back to env-or-nalgebra: just check
        // it builds a working backend.
        let mut b = make_linear_backend(None).unwrap();
        b.set_problem(1, 1).unwrap();
        b.set_linearisation(&[(0, 0, 2.0)]).unwrap();
        let mut rhs = [4.0];
        b.solve_in_place(&mut rhs).unwrap();
        assert!((rhs[0] - 2.0).abs() < 1e-9);
    }
}
