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,
    pub max_zero_cos_iters: 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,
            max_zero_cos_iters: 2,
        }
    }
}

#[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,
    ResidualConverged,
}

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> {
    pub a: &'a mut [T],
    pub b: &'a mut [T],
    pub c: &'a mut [T],
    pub d: &'a mut [T],
    pub e: &'a mut [T],
}

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

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

#[allow(non_snake_case)]
pub trait Minres<T, P>: LinearMap<T, P>
where
    T: Float,
    P: Plugin<T>,
{
    fn minres(
        &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: mut v_k, c: mut v_next, d: mut p_prev, e: mut p_k } = buckets;
        let &Settings { abs_tolerance, rel_tolerance, max_iterations, max_zero_cos_iters } = settings;

        // Initialize auxiliary parameters
        let mut ta: T;
        let mut tb: T = T::zero();

        let mut X: T;
        let mut Y: T = T::zero();

        let mut cos: T = -T::one();
        let mut sin: T = T::zero();

        let mut ua: T;
        let mut ub: T;
        let mut uc: T = T::zero();
        let mut uc_next: T;

        let mut alpha: T;
        let mut residual: T = T::zero();
        let mut scaled_reltol = T::zero();

        let mut constant_residual_iters: usize = 0;

        self.apply(span_extension, x, plugin);

        v_k.iter_mut().zip(rhs.iter().zip(span_extension.iter())).for_each(|(v, (b, Ax))| {
            *v = *b - *Ax;
            residual += *v * *v;
            scaled_reltol += *b * *b;
        });

        residual = residual.sqrt();

        let scaled_reltol = scaled_reltol.sqrt() * rel_tolerance;

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

        v_k.iter_mut().for_each(|v| *v /= residual);

        v_next.reset();
        p_prev.reset();
        p_k.reset();

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

            (ta, tb) = lanczos_step(v_next, v_k, span_extension, tb);

            // Apply Old Householder Reflection
            ub = Y * cos + ta * sin;
            uc_next = tb * sin;

            X = Y * sin - ta * cos;
            Y = -tb * cos;

            // Compute New Householder Reflection
            ua = (X * X + tb * tb).sqrt();
            cos = X / ua;
            sin = tb / ua;

            // Compute Step
            p_k.iter_mut().zip(p_prev.iter().zip(v_k.iter())).for_each(|(curr, (prev, orth))| {
                *curr *= -(uc / ua);
                *curr += -(ub / ua) * *prev + *orth / ua;
            });

            alpha = residual * cos;

            // Update Solution
            x.add(alpha, p_k);
            residual *= sin;

            plugin.peek(kk, x, p_k, alpha, residual);

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

            if sin > T::one() - abs_tolerance {
                constant_residual_iters += 1;

                if constant_residual_iters >= max_zero_cos_iters {
                    return Report::gather(ExitStatus::ResidualConverged, residual, kk);
                }
            } else {
                constant_residual_iters = 0;
            }

            // Close the loop
            core::mem::swap(&mut v_k, &mut v_next);
            core::mem::swap(&mut p_prev, &mut p_k);
            uc = uc_next;
        }

        plugin.end();

        Report::gather(ExitStatus::MaxIterations, residual, max_iterations)
    }
}
