//! A small CSR matrix and a Jacobi-preconditioned BiCGSTAB for the
//! curvilinear patch's pressure equation (`docs/overset_metal_campaign.md`
//! §5.3): the non-orthogonal operator is not symmetric, the patch is a few
//! thousand cells, and `solve_multigrid_pcg` is hard-wired to the
//! five-point Cartesian stencil.

/// Compressed sparse rows.
#[derive(Debug, Clone)]
pub struct CsrMatrix {
    n: usize,
    row_ptr: Vec<usize>,
    col: Vec<usize>,
    val: Vec<f64>,
}

impl CsrMatrix {
    /// Build from `(row, col, value)` triplets; duplicates accumulate,
    /// columns are sorted within each row.
    pub fn from_triplets(n: usize, triplets: &[(usize, usize, f64)]) -> Self {
        let mut rows: Vec<Vec<(usize, f64)>> = vec![Vec::new(); n];
        for &(r, c, v) in triplets {
            rows[r].push((c, v));
        }
        let mut row_ptr = Vec::with_capacity(n + 1);
        let mut col = Vec::with_capacity(triplets.len());
        let mut val = Vec::with_capacity(triplets.len());
        row_ptr.push(0);
        for row in rows.iter_mut() {
            row.sort_by_key(|e| e.0);
            let mut last: Option<usize> = None;
            for &(c, v) in row.iter() {
                if last == Some(c) {
                    *val.last_mut().expect("entry") += v;
                } else {
                    col.push(c);
                    val.push(v);
                    last = Some(c);
                }
            }
            row_ptr.push(col.len());
        }
        Self {
            n,
            row_ptr,
            col,
            val,
        }
    }

    /// Dimension.
    pub fn n(&self) -> usize {
        self.n
    }
    /// Stored entries.
    pub fn nnz(&self) -> usize {
        self.val.len()
    }
    /// `y = A x`.
    pub fn matvec(&self, x: &[f64], y: &mut [f64]) {
        for r in 0..self.n {
            let mut acc = 0.0;
            for k in self.row_ptr[r]..self.row_ptr[r + 1] {
                acc += self.val[k] * x[self.col[k]];
            }
            y[r] = acc;
        }
    }
    /// The diagonal (zero where absent).
    pub fn diagonal(&self) -> Vec<f64> {
        let mut d = vec![0.0; self.n];
        for r in 0..self.n {
            for k in self.row_ptr[r]..self.row_ptr[r + 1] {
                if self.col[k] == r {
                    d[r] = self.val[k];
                }
            }
        }
        d
    }
    /// Replace row `r` by the identity row (`x_r = b_r`): the anchor of a
    /// pure-Neumann problem.
    pub fn set_row_identity(&mut self, r: usize) {
        let mut has_diag = false;
        for k in self.row_ptr[r]..self.row_ptr[r + 1] {
            self.val[k] = if self.col[k] == r {
                has_diag = true;
                1.0
            } else {
                0.0
            };
        }
        assert!(has_diag, "row {r} has no diagonal entry to anchor");
    }
    /// Row `r` as `(columns, values)`.
    pub fn row(&self, r: usize) -> (&[usize], &[f64]) {
        let (a, b) = (self.row_ptr[r], self.row_ptr[r + 1]);
        (&self.col[a..b], &self.val[a..b])
    }
}

/// Outcome of a BiCGSTAB solve.
#[derive(Debug, Clone, Copy)]
#[must_use]
pub struct BicgstabResult {
    /// Iterations taken.
    pub iterations: usize,
    /// L1 norm of the true residual `b − A x` at exit.
    pub residual: f64,
    /// Whether the residual reached the tolerance.
    pub converged: bool,
}

/// Subtract the mean from `v` (the consistency projection for a singular
/// right-hand side).
pub fn project_mean(v: &mut [f64]) {
    let mean = v.iter().sum::<f64>() / v.len() as f64;
    for x in v.iter_mut() {
        *x -= mean;
    }
}

/// Jacobi-preconditioned BiCGSTAB (van der Vorst 1992) with an L1
/// true-residual stop; `x` is the initial guess and the result.
pub fn bicgstab_jacobi(
    a: &CsrMatrix,
    b: &[f64],
    x: &mut [f64],
    tolerance: f64,
    max_iterations: usize,
) -> BicgstabResult {
    let n = a.n();
    let diag = a.diagonal();
    let inv_diag: Vec<f64> = diag
        .iter()
        .map(|&d| if d != 0.0 { 1.0 / d } else { 1.0 })
        .collect();
    let l1 = |v: &[f64]| v.iter().map(|t| t.abs()).sum::<f64>();
    let dot = |u: &[f64], v: &[f64]| u.iter().zip(v).map(|(p, q)| p * q).sum::<f64>();

    let mut r = vec![0.0; n];
    a.matvec(x, &mut r);
    for i in 0..n {
        r[i] = b[i] - r[i];
    }
    let mut res = l1(&r);
    if res <= tolerance {
        return BicgstabResult {
            iterations: 0,
            residual: res,
            converged: true,
        };
    }
    let r0 = r.clone();
    let mut p = vec![0.0; n];
    let mut v = vec![0.0; n];
    let mut s = vec![0.0; n];
    let mut t = vec![0.0; n];
    let mut y = vec![0.0; n];
    let mut z = vec![0.0; n];
    let (mut rho_old, mut alpha, mut omega) = (1.0, 1.0, 1.0);

    for it in 1..=max_iterations {
        let rho = dot(&r0, &r);
        if rho == 0.0 || !rho.is_finite() {
            break;
        }
        let beta = (rho / rho_old) * (alpha / omega);
        for i in 0..n {
            p[i] = r[i] + beta * (p[i] - omega * v[i]);
        }
        for i in 0..n {
            y[i] = inv_diag[i] * p[i];
        }
        a.matvec(&y, &mut v);
        let r0v = dot(&r0, &v);
        if r0v == 0.0 || !r0v.is_finite() {
            break;
        }
        alpha = rho / r0v;
        for i in 0..n {
            s[i] = r[i] - alpha * v[i];
        }
        if l1(&s) <= tolerance {
            for i in 0..n {
                x[i] += alpha * y[i];
            }
            a.matvec(x, &mut r);
            for i in 0..n {
                r[i] = b[i] - r[i];
            }
            res = l1(&r);
            return BicgstabResult {
                iterations: it,
                residual: res,
                converged: res <= tolerance,
            };
        }
        for i in 0..n {
            z[i] = inv_diag[i] * s[i];
        }
        a.matvec(&z, &mut t);
        let tt = dot(&t, &t);
        omega = if tt > 0.0 { dot(&t, &s) / tt } else { 0.0 };
        for i in 0..n {
            x[i] += alpha * y[i] + omega * z[i];
            r[i] = s[i] - omega * t[i];
        }
        rho_old = rho;
        // The recurrence residual drifts from the true one; check the true
        // residual whenever the recurrence claims convergence.
        if l1(&r) <= tolerance {
            a.matvec(x, &mut r);
            for i in 0..n {
                r[i] = b[i] - r[i];
            }
            res = l1(&r);
            if res <= tolerance {
                return BicgstabResult {
                    iterations: it,
                    residual: res,
                    converged: true,
                };
            }
        }
        if omega == 0.0 {
            break;
        }
    }
    a.matvec(x, &mut r);
    for i in 0..n {
        r[i] = b[i] - r[i];
    }
    res = l1(&r);
    BicgstabResult {
        iterations: max_iterations,
        residual: res,
        converged: res <= tolerance,
    }
}

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

    #[test]
    fn a_small_nonsymmetric_system_is_solved_to_rounding() {
        // Diagonally dominant, non-symmetric.
        let dense = [
            [4.0, -1.0, 0.0, 0.5, 0.0],
            [-0.5, 5.0, -1.0, 0.0, 0.2],
            [0.0, -1.5, 6.0, -1.0, 0.0],
            [0.1, 0.0, -1.0, 4.0, -1.0],
            [0.0, 0.3, 0.0, -0.5, 3.0],
        ];
        let mut tri = Vec::new();
        for (r, row) in dense.iter().enumerate() {
            for (c, &v) in row.iter().enumerate() {
                if v != 0.0 {
                    tri.push((r, c, v));
                }
            }
        }
        // Duplicate entry: must accumulate.
        tri.push((0, 1, -0.5));
        tri.push((0, 1, 0.5));
        let a = CsrMatrix::from_triplets(5, &tri);
        assert_eq!(a.nnz(), 17);
        let x_true = [1.0, -2.0, 3.0, 0.5, -1.5];
        let mut b = vec![0.0; 5];
        a.matvec(&x_true, &mut b);
        let mut x = vec![0.0; 5];
        let out = bicgstab_jacobi(&a, &b, &mut x, 1e-13, 100);
        assert!(out.converged, "{out:?}");
        for i in 0..5 {
            assert!(
                (x[i] - x_true[i]).abs() < 1e-11,
                "x[{i}] = {} vs {}",
                x[i],
                x_true[i]
            );
        }
    }

    #[test]
    fn an_anchored_periodic_laplacian_recovers_a_periodic_field_up_to_a_constant() {
        // 1-D periodic second difference (singular): anchor one row, project
        // the mean out of the right-hand side.
        let n = 64;
        let h = 1.0 / n as f64;
        let mut tri = Vec::new();
        for i in 0..n {
            tri.push((i, i, 2.0 / (h * h)));
            tri.push((i, (i + 1) % n, -1.0 / (h * h)));
            tri.push((i, (i + n - 1) % n, -1.0 / (h * h)));
        }
        let mut a = CsrMatrix::from_triplets(n, &tri);
        let phi = |x: f64| (2.0 * std::f64::consts::PI * x).sin();
        let exact: Vec<f64> = (0..n).map(|i| phi((i as f64 + 0.5) * h)).collect();
        let mut b = vec![0.0; n];
        a.matvec(&exact, &mut b);
        b[3] += 1e-3; // an inconsistent perturbation the projection must remove
        project_mean(&mut b);
        a.set_row_identity(0);
        b[0] = 0.0;
        let mut x = vec![0.0; n];
        // 1e-9 absolute: |A| ~ 1/h² = 4e3 and |x| ~ 1 put the rounding floor near 1e-10.
        let out = bicgstab_jacobi(&a, &b, &mut x, 1e-9, 2000);
        assert!(out.converged, "{out:?}");
        let shift = exact[0] - x[0];
        let worst = (0..n)
            .map(|i| (x[i] + shift - exact[i]).abs())
            .fold(0.0, f64::max);
        assert!(worst < 2e-3, "worst {worst:.3e}"); // the 1e-3 perturbation's response bounds it
    }
}