mdarray-linalg 0.2.0

use std::ops::{Add, Mul};

use mdarray::{Array, Dim, Layout, Shape, Slice};
use num_complex::ComplexFloat;
use num_traits::Zero;

use super::simple::naive_outer;
use crate::{
    Naive,
    matvec::{Argmax, MatVec, MatVecBuilder, Outer, OuterBuilder, VecOps},
    utils::unravel_index,
};

struct NaiveMatVecBuilder<'a, T, La, Lx, D0, D1>
where
    La: Layout,
    Lx: Layout,
    D0: Dim,
    D1: Dim,
{
    alpha: T,
    a: &'a Slice<T, (D0, D1), La>,
    x: &'a Slice<T, (D1,), Lx>,
}

impl<'a, T, La, Lx, D0, D1> MatVecBuilder<'a, T, La, Lx, D0, D1>
    for NaiveMatVecBuilder<'a, T, La, Lx, D0, D1>
where
    La: Layout,
    Lx: Layout,
    T: ComplexFloat,
    D0: Dim,
    D1: Dim,
{
    /// `α := α·α'`
    fn scale(mut self, alpha: T) -> Self {
        self.alpha = alpha * self.alpha;
        self
    }

    fn eval(self) -> Array<T, (D0,)> {
        let ash = *self.a.shape();
        let (m, n) = (ash.dim(0), ash.dim(1));
        let x_len = self.x.shape().dim(0);

        assert!(n == x_len, "Matrix columns must match vector length");

        let result_shape = <(D0,) as Shape>::from_dims(&[m]);
        let mut result = Array::<T, (D0,)>::from_elem(result_shape, T::zero());

        for i in 0..m {
            let mut sum = T::zero();
            for j in 0..n {
                sum = sum + self.a[[i, j]] * self.x[[j]];
            }
            result[[i]] = self.alpha * sum;
        }
        result
    }

    fn write<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>) {
        let ash = *self.a.shape();
        let (m, n) = (ash.dim(0), ash.dim(1));
        let x_len = self.x.shape().dim(0);
        let y_len = y.shape().dim(0);

        assert!(n == x_len, "Matrix columns must match vector length");
        assert!(m == y_len, "Matrix rows must match y vector length");

        for i in 0..m {
            let mut sum = T::zero();
            for j in 0..n {
                sum = sum + self.a[[i, j]] * self.x[[j]];
            }
            y[[i]] = self.alpha * sum;
        }
    }

    fn add_to_vec<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>) {
        let ash = *self.a.shape();
        let (m, n) = (ash.dim(0), ash.dim(1));
        let x_len = self.x.shape().dim(0);
        let y_len = y.shape().dim(0);

        assert!(n == x_len, "Matrix columns must match x vector length");
        assert!(m == y_len, "Matrix rows must match y vector length");

        for i in 0..m {
            for j in 0..n {
                y[[i]] = y[[i]] + self.alpha * self.a[[i, j]] * self.x[[j]];
            }
        }
    }

    fn add_to_scaled_vec<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>, beta: T) {
        let ash = *self.a.shape();
        let (m, n) = (ash.dim(0), ash.dim(1));
        let x_len = self.x.shape().dim(0);
        let y_len = y.shape().dim(0);

        assert!(n == x_len, "Matrix columns must match x vector length");
        assert!(m == y_len, "Matrix rows must match y vector length");

        for i in 0..m {
            y[[i]] = beta * y[[i]];
        }

        for i in 0..m {
            for j in 0..n {
                y[[i]] = y[[i]] + self.alpha * self.a[[i, j]] * self.x[[j]];
            }
        }
    }
}

impl<T, D0: Dim, D1: Dim> MatVec<T, D0, D1> for Naive
where
    T: ComplexFloat,
{
    fn matvec<'a, La, Lx>(
        &self,
        a: &'a Slice<T, (D0, D1), La>,
        x: &'a Slice<T, (D1,), Lx>,
    ) -> impl MatVecBuilder<'a, T, La, Lx, D0, D1>
    where
        La: Layout,
        Lx: Layout,
    {
        NaiveMatVecBuilder {
            alpha: T::one(),
            a,
            x,
        }
    }
}

impl<T: ComplexFloat + Add<Output = T> + Mul<Output = T> + Zero + Copy, D: Dim> VecOps<T, D>
    for Naive
{
    type Real = T::Real;

    fn add_to_scaled<Lx: Layout, Ly: Layout>(
        &self,
        alpha: T,
        x: &Slice<T, (D,), Lx>,
        y: &mut Slice<T, (D,), Ly>,
    ) {
        for (elem_x, elem_y) in std::iter::zip(x, y) {
            *elem_y = alpha * (*elem_x) + *elem_y;
        }
    }

    fn dot<Lx: Layout, Ly: Layout>(&self, x: &Slice<T, (D,), Lx>, y: &Slice<T, (D,), Ly>) -> T {
        let mut result = T::zero();
        for (elem_x, elem_y) in std::iter::zip(x, y) {
            result = result + *elem_x * (*elem_y);
        }
        result
    }

    fn dotc<Lx: Layout, Ly: Layout>(&self, x: &Slice<T, (D,), Lx>, y: &Slice<T, (D,), Ly>) -> T {
        let mut result = T::zero();
        for (elem_x, elem_y) in std::iter::zip(x, y) {
            result = result + elem_x.conj() * (*elem_y);
        }
        result
    }

    fn norm2<Lx: Layout>(&self, x: &Slice<T, (D,), Lx>) -> Self::Real {
        let mut sum_sq = T::Real::zero();
        for elem in x.into_iter() {
            sum_sq = sum_sq + elem.abs().powi(2);
        }
        sum_sq.sqrt()
    }

    fn norm1<Lx: Layout>(&self, x: &Slice<T, (D,), Lx>) -> Self::Real {
        let mut sum = T::Real::zero();
        for elem in x.into_iter() {
            sum = sum + elem.re().abs() + elem.im().abs();
        }
        sum
    }

    fn rot<Lx: Layout, Ly: Layout>(
        &self,
        x: &mut Slice<T, (D,), Lx>,
        y: &mut Slice<T, (D,), Ly>,
        c: Self::Real,
        s: T,
    ) {
        // Apply a Givens rotation to vectors x and y in-place:
        //   x[i] =  c * x[i] + s  * y[i]
        //   y[i] = -s* * x[i] + c * y[i]
        // where c is real, s is (possibly complex), and c² + |s|² = 1.
        //
        // For real types, s* == s and this reduces to the standard real Givens rotation.
        // This matches the BLAS `drot`/`zrot` convention.
        for (elem_x, elem_y) in std::iter::zip(x, y) {
            // Store original x before overwriting it
            let old_x = *elem_x;
            let old_y = *elem_y;

            // c is T::Real, so we need to cast it to T for arithmetic with s
            let c_as_t = T::from(c).unwrap();

            *elem_x = c_as_t * old_x + s * old_y;
            *elem_y = c_as_t * old_y - s.conj() * old_x;
        }
    }
}

impl<T: ComplexFloat<Real = T> + PartialOrd + Add<Output = T> + Mul<Output = T> + Zero + Copy>
    Argmax<T> for Naive
{
    fn argmax_write<Lx: Layout, S: Shape>(
        &self,
        x: &Slice<T, S, Lx>,
        output: &mut Vec<usize>,
    ) -> bool {
        output.clear();

        if x.is_empty() {
            return false;
        }

        if x.rank() == 0 {
            return true;
        }

        let mut max_flat_idx = 0;
        let mut max_val = x.iter().next().unwrap();

        for (flat_idx, val) in x.iter().enumerate().skip(1) {
            if val > max_val {
                max_val = val;
                max_flat_idx = flat_idx;
            }
        }

        let indices = unravel_index(x, max_flat_idx);
        output.extend_from_slice(&indices);
        true
    }

    fn argmax<Lx: Layout, S: Shape>(&self, x: &Slice<T, S, Lx>) -> Option<Vec<usize>> {
        let mut result = Vec::new();
        if self.argmax_write(x, &mut result) {
            Some(result)
        } else {
            None
        }
    }

    fn argmax_abs_write<Lx: Layout, S: Shape>(
        &self,
        x: &Slice<T, S, Lx>,
        output: &mut Vec<usize>,
    ) -> bool {
        output.clear();

        if x.is_empty() {
            return false;
        }

        if x.rank() == 0 {
            return true;
        }

        let mut max_flat_idx = 0;
        let mut max_val = x.iter().next().unwrap().abs();

        for (flat_idx, val) in x.iter().enumerate().skip(1) {
            if val.abs() > max_val {
                max_val = val.abs();
                max_flat_idx = flat_idx;
            }
        }

        let indices = unravel_index(x, max_flat_idx);
        output.extend_from_slice(&indices);
        true
    }

    fn argmax_abs<Lx: Layout, S: Shape>(&self, x: &Slice<T, S, Lx>) -> Option<Vec<usize>> {
        let mut result = Vec::new();
        if self.argmax_abs_write(x, &mut result) {
            Some(result)
        } else {
            None
        }
    }
}

impl<T, Dx, Dy> Outer<T, Dx, Dy> for Naive
where
    T: ComplexFloat,
    Dx: Dim,
    Dy: Dim,
{
    fn outer<'a, Lx, Ly>(
        &self,
        x: &'a Slice<T, (Dx,), Lx>,
        y: &'a Slice<T, (Dy,), Ly>,
    ) -> impl OuterBuilder<'a, T, Lx, Ly, Dx, Dy>
    where
        Lx: Layout,
        Ly: Layout,
    {
        NaiveOuterBuilder {
            alpha: T::one(),
            x,
            y,
        }
    }
}

struct NaiveOuterBuilder<'a, T, Lx, Ly, Dx, Dy>
where
    Lx: Layout,
    Ly: Layout,
    Dx: Dim,
    Dy: Dim,
{
    alpha: T,
    x: &'a Slice<T, (Dx,), Lx>,
    y: &'a Slice<T, (Dy,), Ly>,
}

impl<'a, T, Lx, Ly, Dx, Dy> OuterBuilder<'a, T, Lx, Ly, Dx, Dy>
    for NaiveOuterBuilder<'a, T, Lx, Ly, Dx, Dy>
where
    Lx: Layout,
    Ly: Layout,
    Dx: Dim,
    Dy: Dim,
    T: ComplexFloat,
{
    /// `α := α·α'`
    fn scale(mut self, alpha: T) -> Self {
        self.alpha = alpha * self.alpha;
        self
    }

    /// Returns `α·xy`
    fn eval(self) -> Array<T, (Dx, Dy)> {
        let m = self.x.shape().dim(0);
        let n = self.y.shape().dim(0);

        let a_shape = <(Dx, Dy) as Shape>::from_dims(&[m, n]);
        let mut a = Array::<T, (Dx, Dy)>::from_elem(a_shape, T::zero());

        naive_outer(&mut a, self.x, self.y, self.alpha);

        a
    }

    /// `a := α·xy`
    fn write<La: Layout>(self, a: &mut Slice<T, (Dx, Dy), La>) {
        let m = self.x.shape().dim(0);
        let n = self.y.shape().dim(0);

        let ash = *a.shape();
        let (ma, na) = (ash.dim(0), ash.dim(1));

        assert!(ma == m, "Output shape must match input vector length");
        assert!(na == n, "Output shape must match input vector length");

        naive_outer(a, self.x, self.y, self.alpha);
    }

    /// Rank-1 update: `A := α·x·yᵀ + A`
    fn add_to<La: Layout>(self, a: &mut Slice<T, (Dx, Dy), La>) {
        let m = self.x.shape().dim(0);
        let n = self.y.shape().dim(0);

        let ash = *a.shape();
        let (ma, na) = (ash.dim(0), ash.dim(1));

        assert!(ma == m, "Output shape must match input vector length");
        assert!(na == n, "Output shape must match input vector length");

        naive_outer(a, self.x, self.y, self.alpha);
    }

}