pub use crate::algebra::*;
pub use crate::float::*;
pub use crate::plugins::*;

#[derive(Debug, Clone, Copy)]
pub struct Settings<T: Float> {
    pub abs_tolerance: T,
    pub rel_tolerance: T,
    pub max_iterations: usize,
}

impl<T: Float> Default for Settings<T> {
    fn default() -> Self {
        Self { abs_tolerance: T::epsilon().sqrt(), rel_tolerance: T::epsilon().sqrt(), max_iterations: 1_000 }
    }
}

#[derive(Debug, Clone)]
pub struct Report<T> {
    pub status: ExitStatus,
    pub residual: T,
    pub iterations: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExitStatus {
    RelativeSuccess,
    AbsoluteSuccess,
    MaxIterations,
}

impl<T> Report<T> {
    pub fn gather(status: ExitStatus, residual: T, iterations: usize) -> Self {
        Self { status, residual, iterations }
    }
}

// ---

#[derive(Debug)]
pub struct Buckets<'a, T> {
    a: &'a mut [T],
    b: &'a mut [T],
    c: &'a mut [T],
}

impl<'a, T> Buckets<'a, T> {
    pub fn new(a: &'a mut [T], b: &'a mut [T], c: &'a mut [T]) -> Self {
        Self { a, b, c }
    }
}

impl<F, P, T> ConjugateGradient<F, P> for T
where
    F: Float,
    P: Plugin<F>,
    T: LinearMap<F, P>,
{
}

#[allow(non_snake_case)]
pub trait ConjugateGradient<T, P>: LinearMap<T, P>
where
    T: Float,
    P: Plugin<T>,
{
    fn cg(
        &self,
        x: &mut [T], // Solution vector (initial guess)
        rhs: &[T],   // Right-hand-side of the system
        settings: &Settings<T>,
        plugin: &mut P, // User-defined plugin
        buckets: Buckets<'_, T>,
    ) -> Report<T> {
        plugin.start();

        // Unpack input structs
        let Buckets { a: span_extension, b: r_k, c: p_k } = buckets;
        let &Settings { abs_tolerance, rel_tolerance, max_iterations } = settings;

        // Initialize auxiliary parameters

        let mut alpha: T;
        let mut residual_squared: T = T::zero();
        let mut rs_new: T;
        let mut rhs_norm_squared = T::zero();

        self.apply(span_extension, x, plugin);

        r_k.iter_mut().zip(rhs.iter().zip(span_extension.iter())).for_each(|(r, (b, Ax))| {
            *r = *b - *Ax;
            residual_squared += *r * *r;
            rhs_norm_squared += *b * *b;
        });

        let scaled_reltol = rhs_norm_squared.sqrt() * rel_tolerance;

        if residual_squared.sqrt() <= abs_tolerance {
            plugin.end();
            return Report::gather(ExitStatus::AbsoluteSuccess, residual_squared.sqrt(), 0);
        } else if residual_squared.sqrt() <= scaled_reltol {
            plugin.end();
            return Report::gather(ExitStatus::RelativeSuccess, residual_squared.sqrt(), 0);
        }

        p_k.copy_from_slice(r_k);

        for kk in 1..=max_iterations {
            self.apply(span_extension, p_k, plugin);

            alpha = residual_squared / p_k.dot(span_extension);

            plugin.peek(kk, x, p_k, alpha, residual_squared.sqrt());

            rs_new = T::zero();
            x.iter_mut().zip(p_k.iter_mut()).zip(r_k.iter_mut().zip(span_extension.iter())).for_each(
                |((x, p), (r, ap))| {
                    *x += alpha * *p;
                    *r -= alpha * *ap;
                    rs_new += *r * *r;
                },
            );

            p_k.iter_mut().zip(r_k.iter()).for_each(|(p, r)| {
                *p *= rs_new / residual_squared;
                *p += *r;
            });

            residual_squared = rs_new;

            if residual_squared.sqrt() <= abs_tolerance {
                plugin.end();
                return Report::gather(ExitStatus::AbsoluteSuccess, residual_squared.sqrt(), kk);
            } else if residual_squared.sqrt() <= scaled_reltol {
                plugin.end();
                return Report::gather(ExitStatus::RelativeSuccess, residual_squared.sqrt(), kk);
            }
        }

        plugin.end();

        Report::gather(ExitStatus::MaxIterations, residual_squared.sqrt(), max_iterations)
    }
}
