use core::fmt::{Debug, Display, LowerExp, UpperExp};
use core::iter::{Product, Sum};
use core::ops::{AddAssign, DivAssign, MulAssign, RemAssign, SubAssign};
use num_traits::float::Float as NumFloat;

/// Trait for generic declarations of floating point types (`f32` and `f64`).
pub trait Float:
    NumFloat
    + AddAssign
    + SubAssign
    + MulAssign
    + DivAssign
    + RemAssign
    + Sum
    + Product
    + LowerExp
    + UpperExp
    + Debug
    + Display
    + Default
{
}

impl Float for f64 {}
impl Float for f32 {}

// ---

pub trait SliceOps<T> {
    fn norm_1(&self) -> T;
    fn norm_2(&self) -> T;
    fn norm_inf(&self) -> T;
    fn max(&self) -> T;

    fn reset(&mut self) -> &mut Self;
    fn scale(&mut self, beta: T) -> &mut Self;

    fn add(&mut self, beta: T, y: &[T]) -> &mut Self;
    fn scale_add(&mut self, alpha: T, beta: T, y: &[T]) -> &mut Self;
    fn linear_comb(&mut self, a: T, x: &[T], b: T, y: &[T]) -> &mut Self;

    fn take_eval(&mut self, x: &[T], y: &[T], f: impl Fn(T, T) -> T) -> &mut Self;
    fn diag_scale(&mut self, x: &[T]) -> &mut Self;

    fn dot(&self, other: &Self) -> T;
}

impl<T: Float> SliceOps<T> for [T] {
    #[inline(always)]
    fn norm_1(&self) -> T {
        self.iter().map(|x| x.abs()).sum::<T>()
    }

    #[inline(always)]
    fn norm_2(&self) -> T {
        self.iter().map(|x| x.powi(2)).sum::<T>().sqrt()
    }

    #[inline(always)]
    fn norm_inf(&self) -> T {
        self.iter().fold(T::zero(), |a, &b| a.max(b.abs()))
    }

    #[inline(always)]
    fn max(&self) -> T {
        self.iter().fold(T::neg_infinity(), |a, &b| a.max(b))
    }

    #[inline(always)]
    fn reset(&mut self) -> &mut Self {
        for elm in self.iter_mut() {
            *elm = T::zero();
        }
        self
    }

    #[inline(always)]
    fn scale(&mut self, beta: T) -> &mut Self {
        for elm in self.iter_mut() {
            *elm *= beta;
        }
        self
    }

    #[inline(always)]
    fn add(&mut self, beta: T, y: &[T]) -> &mut Self {
        for (kk, elm) in self.iter_mut().enumerate() {
            *elm += beta * y[kk];
        }
        self
    }

    #[inline(always)]
    fn scale_add(&mut self, alpha: T, beta: T, y: &[T]) -> &mut Self {
        for (kk, elm) in self.iter_mut().enumerate() {
            *elm *= alpha;
            *elm += beta * y[kk];
        }
        self
    }

    /// Linear combination of two vectors.
    /// The original contents of the mutable input vector are disregarded and overwritten.
    #[inline(always)]
    fn linear_comb(&mut self, alpha: T, x: &[T], beta: T, y: &[T]) -> &mut Self {
        for (kk, elm) in self.iter_mut().enumerate() {
            *elm = alpha * x[kk] + beta * y[kk];
        }
        self
    }

    /// Computes the dot product between two vectors.
    /// Vectors should have the same number of elements.
    #[inline(always)]
    fn dot(&self, other: &Self) -> T {
        self.iter().zip(other.iter()).map(|(x, y)| *x * *y).sum::<T>()
    }

    #[inline(always)]
    fn take_eval(&mut self, x: &[T], y: &[T], f: impl Fn(T, T) -> T) -> &mut Self {
        self.iter_mut().zip(x.iter().zip(y.iter())).for_each(|(r, (a, b))| *r = f(*a, *b));

        self
    }

    fn diag_scale(&mut self, x: &[T]) -> &mut Self {
        self.iter_mut().zip(x.iter()).for_each(|(r, a)| *r *= *a);

        self
    }
}

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

    #[rstest]
    fn test_dot_orthogonal_vectors() {
        let sz = 100;
        let mut v1: Vec<f64> = vec![0.0; sz];
        let mut v2: Vec<f64> = vec![0.0; sz];

        for (kk, val) in v1.iter_mut().enumerate() {
            if kk % 2 == 0 {
                *val = 1.0;
            }
        }

        v2.clone_from_slice(&v1);
        v2.rotate_right(1);

        // Check that vectors are non-zero
        assert!(v1.norm_2() > 0.0);
        assert!(v2.norm_2() > 0.0);

        // Check that the dot product is zero.
        assert_eq!(v1.dot(&v2), 0.0);
    }

    #[rstest]
    fn test_dot_sum_of_integers() {
        let sz: usize = 100;

        // Vector of ones
        let v1: Vec<f64> = vec![1.0; sz];

        // Vector of increasing integers values
        let mut v2: Vec<f64> = vec![0.0; sz];

        for (kk, val) in v2.iter_mut().enumerate() {
            *val = kk as f64 + 1.0;
        }

        // Dot product equals the sum of elements of v2
        assert_eq!(v1.dot(&v2), v2.iter().sum());
    }

    #[rstest]
    fn test_dot_norm() {
        let sz: usize = 100;
        let v1: Vec<f64> = vec![-2.0; sz];

        // Check that the dot product matches the 2-Norm squared.
        assert_eq!(v1.dot(&v1), v1.norm_2().powi(2));
    }

    #[rstest]
    fn take_eval() {
        const N: usize = 999;

        let x = vec![1.0; N];
        let y = vec![4.0; N];

        let a = 3.0;
        let b = -2.0;

        let mut comb = vec![0.0; N];

        comb.take_eval(&x, &y, |l, r| a * l + b * r);
        assert!(comb == vec![-5.0; N]);

        comb.take_eval(&x, &y, |l, r| l * r * 2.0);
        assert!(comb == vec![8.0; N]);
    }
}
