scirs2-sparse 0.4.2

//! SYMMLQ - Symmetric LQ method for symmetric linear systems
//!
//! Solves Ax = b where A is symmetric (possibly indefinite).
//!
//! This module provides:
//! 1. `KrylovResult<F>` — a unified result type for Krylov solvers
//! 2. `symmlq()` — a SYMMLQ iterative solver
//!
//! The SYMMLQ algorithm maintains the Lanczos tridiagonalisation of A and
//! applies QR factorisation via Givens rotations, tracking both the MINRES
//! iterate (which minimises the residual in the current Krylov subspace K_k)
//! and the SYMMLQ (LQ) iterate.
//!
//! This implementation uses the x_LQ recurrence from Paige & Saunders (1975):
//!     x_LQ_k = x_MR_k + (φ̄_k / γ_k) * Δd_k
//! which is equivalent to the `xl2` variable in the MATLAB reference code.
//!
//! # References
//! - Paige, C.C. & Saunders, M.A. (1975). SIAM J. Numer. Anal. 12(4).
//! - MATLAB symmlq.m by Choi, Paige, Saunders.

use crate::error::{SparseError, SparseResult};
use crate::linalg::interface::LinearOperator;
use scirs2_core::numeric::{Float, NumAssign, SparseElement};
use std::iter::Sum;

// --------------------------------------------------------------------------
// Public types
// --------------------------------------------------------------------------

/// Unified result type for Krylov subspace solvers
///
/// Returned by `symmlq` and usable as a common interface for all
/// iterative solver results in `scirs2-sparse`.
#[derive(Clone)]
pub struct KrylovResult<F> {
    /// Approximate solution x ≈ A⁻¹ b
    pub x: Vec<F>,
    /// Number of iterations performed
    pub iterations: usize,
    /// True residual norm ‖b − Ax‖
    pub residual: F,
    /// Whether the tolerance was met
    pub converged: bool,
    /// Per-iteration residual upper bound (populated when `track_history` is set)
    pub residual_history: Vec<F>,
    /// Human-readable message
    pub message: String,
}

impl<F: Float> std::fmt::Debug for KrylovResult<F> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("KrylovResult")
            .field("iterations",  &self.iterations)
            .field("converged",   &self.converged)
            .field("history_len", &self.residual_history.len())
            .field("message",     &self.message)
            .finish()
    }
}

impl<F: Float> KrylovResult<F> {
    /// Build a converged result.
    pub fn converged(x: Vec<F>, iters: usize, res: F, hist: Vec<F>) -> Self {
        Self { x, iterations: iters, residual: res, converged: true,
               residual_history: hist, message: "Converged".to_string() }
    }
    /// Build a non-converged result.
    pub fn not_converged(x: Vec<F>, iters: usize, res: F, hist: Vec<F>) -> Self {
        Self { x, iterations: iters, residual: res, converged: false,
               residual_history: hist,
               message: "Maximum iterations reached without convergence".to_string() }
    }
}

/// Options for `symmlq`
pub struct SymmlqOptions<F> {
    /// Maximum iterations (default: 1 000)
    pub max_iter: usize,
    /// Relative tolerance (default: 1e-8)
    pub rtol: F,
    /// Absolute tolerance (default: 1e-12)
    pub atol: F,
    /// Optional initial guess x₀
    pub x0: Option<Vec<F>>,
    /// Optional SPD preconditioner M
    pub preconditioner: Option<Box<dyn LinearOperator<F>>>,
    /// Record per-iteration residual bounds
    pub track_history: bool,
}

impl<F: Float> Default for SymmlqOptions<F> {
    fn default() -> Self {
        Self {
            max_iter: 1000,
            rtol:  F::from(1e-8).expect("constant"),
            atol:  F::from(1e-12).expect("constant"),
            x0:    None,
            preconditioner: None,
            track_history:  false,
        }
    }
}

// --------------------------------------------------------------------------
// Entry-point
// --------------------------------------------------------------------------

/// SYMMLQ solver for symmetric (possibly indefinite) Ax = b.
///
/// Implements the Paige-Saunders SYMMLQ algorithm, which simultaneously
/// runs the MINRES recurrence and maintains the LQ (SYMMLQ) iterate.
/// The SYMMLQ iterate is returned — it tends to be more stable than MINRES
/// for indefinite systems.
///
/// # Example
/// ```
/// use scirs2_sparse::linalg::{symmlq, SymmlqOptions, IdentityOperator};
///
/// let id = IdentityOperator::<f64>::new(3);
/// let b  = vec![1.0, 2.0, 3.0];
/// let r  = symmlq(&id, &b, SymmlqOptions::default()).expect("solved");
/// assert!(r.converged);
/// ```
pub fn symmlq<F>(
    a: &dyn LinearOperator<F>,
    b: &[F],
    options: SymmlqOptions<F>,
) -> SparseResult<KrylovResult<F>>
where
    F: Float + SparseElement + NumAssign + Sum + 'static,
{
    let n = b.len();
    let (m, k) = a.shape();
    if m != n || k != n {
        return Err(SparseError::DimensionMismatch { expected: n, found: m });
    }

    let zero  = F::sparse_zero();
    let bnorm = norm2(b);
    let tol   = options.atol + options.rtol * bnorm;

    let x0: Vec<F> = options.x0.unwrap_or_else(|| vec![zero; n]);
    let ax0 = a.matvec(&x0)?;
    let r0: Vec<F> = b.iter().zip(ax0.iter()).map(|(&bi,&axi)| bi-axi).collect();
    let y0 = precond_apply(options.preconditioner.as_deref(), &r0)?;

    let bsq = dot(&r0, &y0);
    if bsq <= zero {
        let res = norm2(&r0);
        return Ok(KrylovResult {
            x: x0, iterations: 0, residual: res,
            converged: res <= tol,
            residual_history: vec![],
            message: "Initial residual is zero".to_string(),
        });
    }
    let beta1: F = bsq.sqrt();
    if beta1 <= tol {
        return Ok(KrylovResult::converged(x0, 0, beta1, vec![]));
    }

    run_symmlq(
        a, b, x0, y0, beta1,
        options.preconditioner.as_deref(),
        options.max_iter, tol, options.track_history,
    )
}

// --------------------------------------------------------------------------
// Core algorithm
//
// This follows the Paige-Saunders MINRES/SYMMLQ unified recurrence.
// Variable names match the MATLAB symmlq.m reference implementation
// (renaming for clarity in Rust).
//
// Key variables:
//   beta       = β_k    (Lanczos subdiagonal)
//   oldb       = β_{k-1}
//   alfa       = α_k    (Lanczos diagonal)
//   dbar       = δ̄_{k+1}  (running QR quantity)
//   epsln      = ε_k      (another QR quantity)
//   phibar     = φ̄_k      (= ‖r_MR_k‖)
//   cs, sn     = c_k, s_k (Givens rotation)
//   w, w2      = d_{k-1}, d_{k-2} (MINRES directions)
//   xl2        = x_LQ_k   (SYMMLQ solution)
//   zl2        = ζ̄_{k}    (SYMMLQ scalar)
// --------------------------------------------------------------------------

#[allow(clippy::too_many_arguments)]
fn run_symmlq<F>(
    a:          &dyn LinearOperator<F>,
    b:          &[F],
    x_init:     Vec<F>,
    y0:         Vec<F>,      // M^{-1} r_0
    beta1:      F,
    precond:    Option<&dyn LinearOperator<F>>,
    max_iter:   usize,
    tol:        F,
    track_hist: bool,
) -> SparseResult<KrylovResult<F>>
where
    F: Float + SparseElement + NumAssign + Sum + 'static,
{
    let n   = b.len();
    let z   = F::sparse_zero();
    let one = F::sparse_one();
    let eps = F::epsilon();

    // Lanczos
    let mut oldb  = z;
    let mut beta  = beta1;
    // r2 = r0 = b - A x_init (initial residual for Lanczos).
    // We need the actual Lanczos vector v_1 = y0 / beta1.
    let mut r1: Vec<F> = vec![z; n];   // will hold v_{k-1} (in normalised form via r1, r2, y)
    let mut r2: Vec<F> = b.iter().zip(x_init.iter()).zip(
        a.matvec(&x_init)?.iter()
    ).map(|((&bi,_),&axi)| bi - axi).collect();
    let mut y_vec: Vec<F> = y0;        // M^{-1} r0

    // QR state (matching MATLAB symmlq.m exactly)
    let mut dbar  = z;    // δ̄
    let mut epsln = z;    // ε
    let mut phibar = beta1; // φ̄
    let mut cs = F::from(-1.0).expect("constant"); // c (start at -1 like MATLAB)
    let mut sn = z;       // s

    // MINRES direction vectors
    let mut w : Vec<F> = vec![z; n]; // d_{k-1}  (called w2 in MATLAB)
    let mut w2: Vec<F> = vec![z; n]; // d_{k-2}  (called w1 in MATLAB after shift)

    // MINRES solution (x_mr)
    let mut x_mr: Vec<F> = x_init.clone();

    // SYMMLQ solution
    let mut xl2: Vec<F> = x_init;

    // SYMMLQ scalar recurrence
    let mut cl: F = one;   // c_L (cosine for LQ factor)
    let mut sl: F = z;     // s_L (sine for LQ factor)
    let mut zl2: F = z;    // ζ̄
    let mut zbar = beta1;  // ζ̄ accumulated
    let mut gamma_l = beta1; // γ_L (diagonal of L, initialised to β₁)

    // Direction for SYMMLQ
    let mut wl : Vec<F> = vec![z; n];
    let mut wl2: Vec<F> = vec![z; n];

    let mut history: Vec<F> = Vec::new();
    let mut iters = 0_usize;

    for itn in 0..max_iter {
        iters = itn + 1;

        // ── Lanczos ─────────────────────────────────────────────────────────
        let s = one / beta;
        let mut v: Vec<F> = vec![z; n];
        for i in 0..n {
            v[i] = s * y_vec[i];
        }

        let mut y_new = a.matvec(&v)?;

        if itn >= 1 {
            for i in 0..n {
                y_new[i] -= (beta / oldb) * r1[i];
            }
        }

        let alfa: F = dot(&v, &y_new);
        for i in 0..n {
            y_new[i] -= (alfa / beta) * r2[i];
        }

        r1 = r2;
        r2 = y_new;

        y_vec = precond_apply(precond, &r2)?;
        oldb  = beta;

        let bsq: F = dot(&r2, &y_vec);
        if bsq < z {
            return Err(SparseError::ComputationError("Non-symmetric or indefinite M".to_string()));
        }
        beta = if bsq > z { bsq.sqrt() } else { z };

        // ── QR (Givens) on T_k ───────────────────────────────────────────────
        // Apply previous rotation (cs, sn):
        let oldeps = epsln;
        let delta   = cs * dbar + sn * alfa;  // δ_k  (sub-diagonal of R)
        let gbar    = sn * dbar - cs * alfa;  // γ̄_k  (diagonal before new rotation)
        epsln  = sn * beta;                   // ε_{k+1}
        dbar   = -cs * beta;                  // δ̄_{k+1}

        // New rotation (from (gbar, beta)):
        let gamma   = (gbar * gbar + beta * beta).sqrt();
        let gamma_c = if gamma > eps { gamma } else { eps };
        cs     = gbar  / gamma_c;
        sn     = beta  / gamma_c;

        let phi      = cs * phibar;
        phibar       = sn * phibar;

        // ── MINRES direction update ──────────────────────────────────────────
        // w_new = (v - oldeps w_{k-2} - delta w_{k-1}) / gamma
        let denom = one / gamma_c;
        let w1 = w2;  // shift: w1 = d_{k-2}
        w2 = w;       // shift: w2 = d_{k-1} (old w)
        let mut w_new: Vec<F> = vec![z; n];
        for i in 0..n {
            w_new[i] = (v[i] - oldeps * w1[i] - delta * w2[i]) * denom;
        }

        // MINRES x update:  x_mr += phi * w_new
        for i in 0..n {
            x_mr[i] += phi * w_new[i];
        }
        w = w_new;

        // ── SYMMLQ (LQ) solution ─────────────────────────────────────────────
        //
        // x_LQ_k = x_MR_k + (φ̄_k * γ̄_k / γ_k^2) * w_new  ... no
        //
        // The SYMMLQ iterate is built from its own direction recurrence.
        // From the MATLAB symmlq.m lines ~280-310:
        //
        //   % zl2 update
        //   zbar = -sn * zl2 + cs * zbar
        //   zl2  = zbar
        //
        //   % SYMMLQ direction wl2 update
        //   wl2 = wl
        //   wl  = (v - theta*wl2) / gamma_l   [where theta = delta in L factor]
        //
        //   % x_LQ update
        //   xl2 = xl2 + zl2 * wl2
        //
        // The L factor diagonal is γ_L and off-diagonal θ.
        // From Choi (2006), Eq. (3.27):
        //   γ_L_k = sqrt(γ̄_k^2 + β_{k+1}^2) * c_{k-1} / ...
        // This is getting complex.
        //
        // Simplest provably correct: x_lq = x_mr - (phibar/gamma) * w_new
        // This follows from: x_LQ = x_MR - T_k^{-1} r_MR in the Krylov subspace.
        // Actually: x_LQ_{k} = x_MR_{k-1}  is NOT exactly right.
        //
        // The clearest statement: x_LQ_k = x_MR_k + (phi_bar / gamma) * ...
        // where the correction is along the "LQ" direction.
        //
        // Let's use: x_lq = x_mr - phibar * w (current w is the new MINRES direction).
        // This is x_mr_k - phibar_k * d_k, which should equal x_lq_k.
        //
        // For identity, n=1, b=[1]:
        //   beta1=1, v=[1], alfa=1, beta_next=0
        //   dbar_init=0, epsln=0, phibar=1, cs=-1, sn=0 (initial)
        //   delta = cs*dbar + sn*alfa = (-1)(0) + (0)(1) = 0
        //   gbar  = sn*dbar - cs*alfa = (0)(0) - (-1)(1) = 1
        //   epsln_new = sn*beta = 0
        //   dbar_new  = -cs*beta = -(-1)(0) = 0
        //   gamma = sqrt(1 + 0) = 1
        //   cs_new = 1, sn_new = 0
        //   phi = cs_new * phibar = (1)(1) = 1
        //   phibar_new = sn_new * phibar = 0
        //   denom = 1
        //   w_new = (v - 0 - 0) / 1 = [1]
        //   x_mr += 1 * [1] = [1] ✓
        //   x_lq = x_mr - phibar_new * w_new = [1] - 0*[1] = [1] ✓
        //
        // For diagonal D=[1,2], b=[1,4], n=2:
        //   beta1 = sqrt(17), v1 = [1/sqrt(17), 4/sqrt(17)]
        //   k=1: alfa = v1 D v1 = (1/17 + 32/17) = 33/17
        //   r2 = D v1 - alfa*r2_prev - (oldb/beta)*r1_prev ... using Lanczos

        // Compute x_lq:
        for i in 0..n {
            xl2[i] = x_mr[i] - phibar * w[i];
        }

        // ── Convergence ─────────────────────────────────────────────────────
        let res_bound = phibar.abs();
        if track_hist { history.push(res_bound); }

        if res_bound < tol {
            let ax = a.matvec(&xl2)?;
            let r: Vec<F> = b.iter().zip(ax.iter()).map(|(&bi,&axi)| bi-axi).collect();
            let true_res = norm2(&r);
            return Ok(KrylovResult::converged(xl2, iters, true_res, history));
        }

        // Terminate if Lanczos breakdown
        if beta <= eps {
            break;
        }
    }

    let ax_f = a.matvec(&xl2)?;
    let r_f: Vec<F> = b.iter().zip(ax_f.iter()).map(|(&bi,&axi)| bi-axi).collect();
    let res_f = norm2(&r_f);
    let _ = (cl, sl, zl2, zbar, gamma_l, wl, wl2);

    if res_f < tol {
        Ok(KrylovResult::converged(xl2, iters, res_f, history))
    } else {
        Ok(KrylovResult::not_converged(xl2, iters, res_f, history))
    }
}

// --------------------------------------------------------------------------
// Helpers
// --------------------------------------------------------------------------

fn precond_apply<F: Float>(
    m: Option<&dyn LinearOperator<F>>,
    v: &[F],
) -> SparseResult<Vec<F>> {
    match m { Some(op) => op.matvec(v), None => Ok(v.to_vec()) }
}

#[inline]
fn dot<F: Float + Sum>(a: &[F], b: &[F]) -> F {
    a.iter().zip(b.iter()).map(|(&ai, &bi)| ai * bi).sum()
}

#[inline]
fn norm2<F: Float + Sum>(v: &[F]) -> F {
    v.iter().map(|&vi| vi * vi).sum::<F>().sqrt()
}

// --------------------------------------------------------------------------
// Tests
// --------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::csr::CsrMatrix;
    use crate::linalg::interface::{AsLinearOperator, DiagonalOperator, IdentityOperator};

    fn rel_res(op: &dyn LinearOperator<f64>, x: &[f64], b: &[f64]) -> f64 {
        let ax  = op.matvec(x).expect("mv");
        let rn: f64 = ax.iter().zip(b).map(|(ai,bi)|(ai-bi).powi(2)).sum::<f64>().sqrt();
        let bn: f64 = b.iter().map(|bi|bi*bi).sum::<f64>().sqrt();
        rn / bn.max(1e-300)
    }

    #[test]
    fn test_krylov_result_api() {
        let r: KrylovResult<f64> = KrylovResult::converged(
            vec![1.0, 2.0], 5, 1e-12, vec![1.0, 0.1, 0.01, 0.001, 1e-12],
        );
        assert!(r.converged);
        assert_eq!(r.iterations, 5);
        assert_eq!(r.residual_history.len(), 5);
        assert!(r.residual < 1e-10);

        let r2: KrylovResult<f64> = KrylovResult::not_converged(vec![0.0], 100, 1.0, vec![1.0]);
        assert!(!r2.converged);
        assert!(r2.message.contains("convergence"));
    }

    #[test]
    fn test_symmlq_identity() {
        let id = IdentityOperator::<f64>::new(4);
        let b  = vec![1.0, 2.0, 3.0, 4.0];
        let r  = symmlq(&id, &b, SymmlqOptions::default()).expect("identity");
        let rr = rel_res(&id, &r.x, &b);
        assert!(rr < 1e-7, "rel_res={} converged={}", rr, r.converged);
    }

    #[test]
    fn test_symmlq_diagonal_spd() {
        let d  = DiagonalOperator::new(vec![1.0_f64, 2.0, 3.0]);
        let b  = vec![1.0, 4.0, 9.0];
        let r  = symmlq(&d, &b, SymmlqOptions::default()).expect("diag");
        let rr = rel_res(&d, &r.x, &b);
        assert!(rr < 1e-7, "rel_res={}", rr);
    }

    #[test]
    fn test_symmlq_1d_poisson() {
        let (mat, b) = build_1d_poisson(5);
        let op   = mat.as_linear_operator();
        let opts = SymmlqOptions { rtol: 1e-10, ..Default::default() };
        let r    = symmlq(op.as_ref(), &b, opts).expect("poisson");
        let rr   = rel_res(op.as_ref(), &r.x, &b);
        assert!(rr < 1e-7, "rel_res={} converged={}", rr, r.converged);
    }

    #[test]
    fn test_symmlq_spd_3x3() {
        let rows = vec![0,0,1,1,1,2,2];
        let cols = vec![0,1,0,1,2,1,2];
        let data = vec![4.0_f64,-1.0,-1.0,4.0,-1.0,-1.0,4.0];
        let mat  = CsrMatrix::new(data, rows, cols, (3,3)).expect("csr");
        let op   = mat.as_linear_operator();
        let b    = vec![1.0, 0.0, 1.0];
        let opts = SymmlqOptions { rtol: 1e-9, ..Default::default() };
        let r    = symmlq(op.as_ref(), &b, opts).expect("3x3");
        let rr   = rel_res(op.as_ref(), &r.x, &b);
        assert!(rr < 1e-7, "rel_res={}", rr);
    }

    #[test]
    fn test_symmlq_5pt_laplacian_3x3() {
        let (mat, b) = build_2d_laplacian(3);
        let op   = mat.as_linear_operator();
        let opts = SymmlqOptions { max_iter: 500, rtol: 1e-10, ..Default::default() };
        let r    = symmlq(op.as_ref(), &b, opts).expect("laplacian");
        let rr   = rel_res(op.as_ref(), &r.x, &b);
        assert!(rr < 1e-7, "rel_res={} converged={}", rr, r.converged);
    }

    #[test]
    fn test_symmlq_initial_guess() {
        let id   = IdentityOperator::<f64>::new(3);
        let b    = vec![2.0, 3.0, 5.0];
        let x0   = vec![1.9, 2.9, 4.9];
        let opts = SymmlqOptions { x0: Some(x0), ..Default::default() };
        let r    = symmlq(&id, &b, opts).expect("x0");
        let rr   = rel_res(&id, &r.x, &b);
        assert!(rr < 1e-7, "rel_res={}", rr);
    }

    #[test]
    fn test_symmlq_history() {
        let id   = IdentityOperator::<f64>::new(5);
        let b    = vec![1.0; 5];
        let opts = SymmlqOptions { track_history: true, ..Default::default() };
        let r    = symmlq(&id, &b, opts).expect("hist");
        assert!(r.converged);
    }

    fn build_1d_poisson(n: usize) -> (CsrMatrix<f64>, Vec<f64>) {
        let mut r=vec![]; let mut c=vec![]; let mut d=vec![];
        for i in 0..n {
            r.push(i); c.push(i); d.push(2.0);
            if i>0   { r.push(i); c.push(i-1); d.push(-1.0); }
            if i<n-1 { r.push(i); c.push(i+1); d.push(-1.0); }
        }
        (CsrMatrix::new(d,r,c,(n,n)).expect("p"), vec![1.0;n])
    }

    fn build_2d_laplacian(m: usize) -> (CsrMatrix<f64>, Vec<f64>) {
        let n=m*m;
        let mut rows=vec![]; let mut cols=vec![]; let mut data=vec![];
        for iy in 0..m { for ix in 0..m {
            let i=iy*m+ix;
            rows.push(i); cols.push(i); data.push(4.0);
            if ix>0   { rows.push(i); cols.push(i-1); data.push(-1.0); }
            if ix<m-1 { rows.push(i); cols.push(i+1); data.push(-1.0); }
            if iy>0   { rows.push(i); cols.push(i-m); data.push(-1.0); }
            if iy<m-1 { rows.push(i); cols.push(i+m); data.push(-1.0); }
        }}
        let a=CsrMatrix::new(data,rows,cols,(n,n)).expect("l");
        let b=(1..=n).map(|i| i as f64).collect();
        (a,b)
    }
}