diffsol-la 0.3.0

use std::cell::RefCell;

use crate::context::broadcast_batch;
use crate::{
    error::LaError, linear_solver::LinearSolver, linear_solver_error, scalar::IndexType, Context,
    FaerContext, FaerScalar, FaerSparseMat, FaerVec, LinearOp, Matrix,
};

use faer::dyn_stack::{MemBuffer, MemStack};
use faer::reborrow::{Reborrow, ReborrowMut};
use faer::sparse::linalg::lu::{factorize_symbolic_lu, LuRef, NumericLu, SymbolicLu};
use faer::Conj;

/// Sparse $LU$ factorization state for a (possibly batched) [FaerSparseMat].
///
/// The symbolic factorization, one numeric factorization per batch and both scratch buffers are
/// allocated once, in [`FaerSparseLU::set_sparsity`].  faer sizes a [`NumericLu`]'s own buffers on
/// its first `factorize_numeric_lu` rather than on construction, so the first
/// [`FaerSparseLU::set_linearisation`] after that still allocates inside faer; later ones reuse
/// the capacity.
struct SparseLu<T: FaerScalar> {
    a: FaerSparseMat<T>,
    symbolic: SymbolicLu<IndexType>,
    numeric: Vec<NumericLu<IndexType, T>>,
    factor_scratch: MemBuffer,
    solve_scratch: RefCell<MemBuffer>,
}

impl<T: FaerScalar> SparseLu<T> {
    fn new(a: FaerSparseMat<T>) -> Self {
        let symbolic = factorize_symbolic_lu(a.data[0].symbolic(), Default::default())
            .expect("Failed to create symbolic LU");
        let par = a.context.par;
        // the solve always runs one column at a time, so this size holds however many batches the
        // right-hand side broadcasts into
        let factor_scratch =
            MemBuffer::new(symbolic.factorize_numeric_lu_scratch::<T>(par, Default::default()));
        let solve_scratch = RefCell::new(MemBuffer::new(
            symbolic
                .solve_in_place_scratch::<T>(1, par)
                .or(symbolic.solve_transpose_in_place_scratch::<T>(1, par)),
        ));
        Self {
            numeric: (0..a.data.len()).map(|_| NumericLu::new()).collect(),
            a,
            symbolic,
            factor_scratch,
            solve_scratch,
        }
    }

    fn nbatch(&self) -> usize {
        self.numeric.len()
    }

    fn factor(&mut self) {
        let par = self.a.context.par;
        for (numeric, block) in self.numeric.iter_mut().zip(self.a.data.iter()) {
            self.symbolic
                .factorize_numeric_lu(
                    numeric,
                    block.rb(),
                    par,
                    MemStack::new(&mut self.factor_scratch),
                    Default::default(),
                )
                .expect("Failed to factorise matrix");
        }
    }

    fn solve_block(&self, batch: usize, rhs: faer::MatMut<'_, T>, par: faer::Par, transpose: bool) {
        let lu = LuRef::new_unchecked(&self.symbolic, &self.numeric[batch]);
        let mut scratch = self.solve_scratch.borrow_mut();
        let stack = MemStack::new(&mut scratch);
        if transpose {
            lu.solve_transpose_in_place_with_conj(Conj::No, rhs, par, stack);
        } else {
            lu.solve_in_place_with_conj(Conj::No, rhs, par, stack);
        }
    }
}

/// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system.
///
/// The low-level `faer::sparse::linalg::lu` API is used rather than
/// `faer::sparse::linalg::solvers` both so that the parallelism from [`FaerContext::par`] is
/// honoured -- the wrapper types hardcode faer's process-global setting -- and so that the
/// factorizations and scratch space can be allocated once in [`Self::set_sparsity`] and reused.
pub struct FaerSparseLU<T>
where
    T: FaerScalar,
{
    lu: Option<SparseLu<T>>,
}

impl<T> Default for FaerSparseLU<T>
where
    T: FaerScalar,
{
    fn default() -> Self {
        Self { lu: None }
    }
}

impl<T: FaerScalar> LinearSolver<FaerSparseMat<T>> for FaerSparseLU<T> {
    fn set_linearisation<
        C: LinearOp<T = T, V = FaerVec<T>, M = FaerSparseMat<T>, C = FaerContext>,
    >(
        &mut self,
        op: &C,
    ) {
        let lu = self.lu.as_mut().expect("Matrix not set");
        op.matrix_inplace(&mut lu.a);
        lu.factor();
    }

    fn solve_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
        self.solve_in_place_impl(x, false)
    }

    fn solve_transpose_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
        self.solve_in_place_impl(x, true)
    }

    fn set_sparsity<C: LinearOp<T = T, V = FaerVec<T>, M = FaerSparseMat<T>, C = FaerContext>>(
        &mut self,
        op: &C,
    ) {
        let ncols = op.ncols();
        let nrows = op.nrows();
        let a = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
        self.lu = Some(SparseLu::new(a));
    }
}

impl<T: FaerScalar> FaerSparseLU<T> {
    fn solve_in_place_impl(&self, x: &mut FaerVec<T>, transpose: bool) -> Result<(), LaError> {
        let Some(lu) = self.lu.as_ref() else {
            return Err(linear_solver_error!(LuNotInitialized));
        };
        x.context
            .assert_broadcastable_into(lu.nbatch(), "sparse_lu_solve");
        let par = x.context.par;
        let nb = x.data.ncols();
        for batch in 0..nb {
            lu.solve_block(
                broadcast_batch(batch, lu.nbatch(), nb),
                x.data.rb_mut().col_mut(batch).as_mat_mut(),
                par,
                transpose,
            );
        }
        Ok(())
    }
}

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

    crate::linear_solver::generate_lu_tests!(
        f64,
        FaerSparseMat<f64>,
        FaerSparseLU<f64>,
        FaerContext::with_nbatch(2)
    );
    crate::linear_solver::generate_lu_tests!(
        f32,
        FaerSparseMat<f32>,
        FaerSparseLU<f32>,
        FaerContext::with_nbatch(2)
    );
}