basin 0.8.0

use ndarray::{Array1, Array2, ArrayBase, Data, DataMut, Dimension};
use rand::{Rng, RngExt};
use rand_distr::{Distribution, StandardNormal};

use super::cl_scaling::{
    cl_scaling_pair, max_feasible_step_component, project_strictly_inside_component,
    BoxAffineScaling,
};
use super::sample::{SampleStandardNormal, SampleUniformBox};
use super::{
    ClampInPlace, ComponentDivAssign, ComponentMaxAssign, ComponentMulAssign, Dot,
    FloorZerosInPlace, MatDiagonal, MatTransposeVec, MatVec, MatrixFromDiagonal, MatrixIdentity,
    NegInPlace, NormInfinity, NormSquared, RankOneUpdate, ScaleInPlace, ScaledAdd, SymmetricEigen,
    SymmetricEigenError, VectorIndex, VectorLen,
};

impl<S, D> ScaledAdd<f64> for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn scaled_add(&mut self, scalar: f64, other: &Self) {
        assert_eq!(self.shape(), other.shape(), "scaled_add: shape mismatch");
        self.zip_mut_with(other, |x, y| *x += scalar * y);
    }
}

impl<S, D> NormSquared for ArrayBase<S, D>
where
    S: Data<Elem = f64>,
    D: Dimension,
{
    fn norm_squared(&self) -> f64 {
        self.iter().map(|x| x * x).sum()
    }
}

impl<S, D> NormInfinity for ArrayBase<S, D>
where
    S: Data<Elem = f64>,
    D: Dimension,
{
    fn norm_infinity(&self) -> f64 {
        self.iter().map(|x| x.abs()).fold(0.0, f64::max)
    }
}

impl<S, D> Dot for ArrayBase<S, D>
where
    S: Data<Elem = f64>,
    D: Dimension,
{
    fn dot(&self, other: &Self) -> f64 {
        assert_eq!(self.shape(), other.shape(), "dot: shape mismatch");
        self.iter().zip(other.iter()).map(|(a, b)| a * b).sum()
    }
}

impl<S, D> NegInPlace for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn neg_in_place(&mut self) {
        self.map_inplace(|x| *x = -*x);
    }
}

impl SampleUniformBox for Array1<f64> {
    fn sample_uniform_box<R: Rng + ?Sized>(lower: &Self, upper: &Self, rng: &mut R) -> Self {
        assert_eq!(
            lower.len(),
            upper.len(),
            "sample_uniform_box: bounds length mismatch"
        );
        Array1::from_shape_fn(lower.len(), |i| rng.random_range(lower[i]..=upper[i]))
    }
}

impl VectorLen for Array1<f64> {
    fn vec_len(&self) -> usize {
        self.len()
    }
}

// Matrix-vector ops (the only `linalg`-tier traits ndarray implements). They
// lift the backend gate on the linear-constraint solvers; the factorization
// ops stay nalgebra/faer-only, since `ndarray-linalg` needs system BLAS/LAPACK
// (breaks the wasm-default tenet).
//
// These iterate rows/columns explicitly rather than calling ndarray's inherent
// `.dot()`: with basin's own [`Dot`] trait in scope, the bare `dot` method name
// makes inference explore ndarray's recursive `Dot`/`Not` bounds and overflow.
// `assert_eq!` on the lengths matches the panic-on-shape-mismatch contract.

impl MatVec<Array1<f64>> for Array2<f64> {
    fn matvec(&self, x: &Array1<f64>) -> Array1<f64> {
        assert_eq!(
            x.len(),
            self.ncols(),
            "matvec: x has length {} but the matrix has {} columns",
            x.len(),
            self.ncols()
        );
        Array1::from_shape_fn(self.nrows(), |i| {
            self.row(i).iter().zip(x.iter()).map(|(a, xj)| a * xj).sum()
        })
    }
}

impl MatTransposeVec<Array1<f64>> for Array2<f64> {
    fn mat_transpose_vec(&self, x: &Array1<f64>) -> Array1<f64> {
        assert_eq!(
            x.len(),
            self.nrows(),
            "mat_transpose_vec: x has length {} but the matrix has {} rows",
            x.len(),
            self.nrows()
        );
        Array1::from_shape_fn(self.ncols(), |j| {
            self.column(j)
                .iter()
                .zip(x.iter())
                .map(|(a, xi)| a * xi)
                .sum()
        })
    }
}

impl MatrixIdentity for Array2<f64> {
    fn identity(n: usize) -> Self {
        Array2::eye(n)
    }
}

impl MatrixFromDiagonal<Array1<f64>> for Array2<f64> {
    fn from_diagonal(diag: &Array1<f64>) -> Self {
        Array2::from_diag(diag)
    }
}

impl MatDiagonal<Array1<f64>> for Array2<f64> {
    fn diagonal(&self) -> Array1<f64> {
        assert_eq!(
            self.nrows(),
            self.ncols(),
            "diagonal: matrix must be square, got {}x{}",
            self.nrows(),
            self.ncols()
        );
        self.diag().to_owned()
    }
}

impl RankOneUpdate<Array1<f64>> for Array2<f64> {
    fn rank_one_update(&mut self, alpha: f64, v: &Array1<f64>) {
        assert_eq!(
            self.nrows(),
            self.ncols(),
            "rank_one_update: matrix must be square, got {}x{}",
            self.nrows(),
            self.ncols()
        );
        assert_eq!(
            self.nrows(),
            v.len(),
            "rank_one_update: matrix is {}x{} but v has length {}",
            self.nrows(),
            self.ncols(),
            v.len()
        );
        // self[i, j] ← self[i, j] + α · v[i] · v[j] — the symmetric `u == v`
        // case of `general_rank_one_update`. Explicit double-loop matches the
        // `DenseMatrix` pattern; ndarray's `outer` / broadcasting forms would
        // allocate.
        let n = self.nrows();
        for i in 0..n {
            let av = alpha * v[i];
            let mut row = self.row_mut(i);
            for j in 0..n {
                row[j] += av * v[j];
            }
        }
    }
}

impl SymmetricEigen<Array1<f64>> for Array2<f64> {
    fn try_eigh(&self) -> Result<(Self, Array1<f64>), SymmetricEigenError> {
        assert_eq!(
            self.nrows(),
            self.ncols(),
            "try_eigh: matrix must be square, got {}x{}",
            self.nrows(),
            self.ncols()
        );
        let n = self.nrows();
        // `jacobi_eigen` takes a row-major `&[f64]`. `as_standard_layout`
        // returns a `CowArray` that is contiguous in row-major (C) order —
        // borrowing when `self` is already standard, otherwise cloning.
        let standard = self.as_standard_layout();
        let slice = standard
            .as_slice()
            .expect("as_standard_layout produces a contiguous row-major slice");
        let (eigenvalues, eigenvectors) =
            super::dense_eig::jacobi_eigen(slice, n).ok_or(SymmetricEigenError::Failed)?;
        // `jacobi_eigen` returns the eigenvectors row-major with column `k`
        // the eigenvector for `eigenvalues[k]` — exactly what
        // `Array2::from_shape_vec` with the default C-order produces.
        let b =
            Array2::from_shape_vec((n, n), eigenvectors).expect("jacobi_eigen returns n*n entries");
        Ok((b, Array1::from_vec(eigenvalues)))
    }
}

impl VectorIndex for Array1<f64> {
    fn get_scalar(&self, i: usize) -> f64 {
        self[i]
    }
    fn set_scalar(&mut self, i: usize, value: f64) {
        self[i] = value;
    }
}

impl SampleStandardNormal for Array1<f64> {
    fn sample_standard_normal<R: Rng + ?Sized>(template: &Self, rng: &mut R) -> Self {
        Array1::from_shape_fn(template.len(), |_| StandardNormal.sample(rng))
    }
}

impl<S, D> ScaleInPlace for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn scale_in_place(&mut self, scalar: f64) {
        self.map_inplace(|x| *x *= scalar);
    }
}

impl<S, D> ComponentMulAssign for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn component_mul_assign(&mut self, other: &Self) {
        assert_eq!(
            self.shape(),
            other.shape(),
            "component_mul_assign: shape mismatch"
        );
        self.zip_mut_with(other, |x, y| *x *= y);
    }
}

impl<S, D> ComponentMaxAssign for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn component_max_assign(&mut self, other: &Self) {
        assert_eq!(
            self.shape(),
            other.shape(),
            "component_max_assign: shape mismatch"
        );
        self.zip_mut_with(other, |x, y| *x = x.max(*y));
    }
}

impl<S, D> FloorZerosInPlace for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn floor_zeros_in_place(&mut self, value: f64) {
        self.map_inplace(|x| {
            if *x <= 0.0 {
                *x = value;
            }
        });
    }
}

impl<S, D> ComponentDivAssign for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn component_div_assign(&mut self, other: &Self) {
        assert_eq!(
            self.shape(),
            other.shape(),
            "component_div_assign: shape mismatch"
        );
        self.zip_mut_with(other, |x, y| *x /= y);
    }
}

impl<S, D> ClampInPlace for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn clamp_in_place(&mut self, lower: &Self, upper: &Self) {
        assert_eq!(
            self.shape(),
            lower.shape(),
            "clamp_in_place: lower shape mismatch"
        );
        assert_eq!(
            self.shape(),
            upper.shape(),
            "clamp_in_place: upper shape mismatch"
        );
        ndarray::Zip::from(self)
            .and(lower)
            .and(upper)
            .for_each(|x, &lo, &hi| *x = x.clamp(lo, hi));
    }
}

impl<S, D> BoxAffineScaling for ArrayBase<S, D>
where
    S: DataMut<Elem = f64>,
    D: Dimension,
{
    fn compute_cl_scaling(
        &self,
        gradient: &Self,
        lower: &Self,
        upper: &Self,
        d_sq: &mut Self,
        c_diag: &mut Self,
    ) {
        assert_eq!(
            self.shape(),
            gradient.shape(),
            "compute_cl_scaling: gradient shape mismatch"
        );
        assert_eq!(
            self.shape(),
            lower.shape(),
            "compute_cl_scaling: lower shape mismatch"
        );
        assert_eq!(
            self.shape(),
            upper.shape(),
            "compute_cl_scaling: upper shape mismatch"
        );
        assert_eq!(
            self.shape(),
            d_sq.shape(),
            "compute_cl_scaling: d_sq shape mismatch"
        );
        assert_eq!(
            self.shape(),
            c_diag.shape(),
            "compute_cl_scaling: c_diag shape mismatch"
        );
        ndarray::Zip::from(d_sq)
            .and(c_diag)
            .and(self)
            .and(gradient)
            .and(lower)
            .and(upper)
            .for_each(|d, c, &x, &g, &l, &u| {
                let (d_sq_i, c_i) = cl_scaling_pair(x, g, l, u);
                *d = d_sq_i;
                *c = c_i;
            });
    }

    fn max_feasible_step(&self, step: &Self, lower: &Self, upper: &Self) -> f64 {
        assert_eq!(
            self.shape(),
            step.shape(),
            "max_feasible_step: step shape mismatch"
        );
        assert_eq!(
            self.shape(),
            lower.shape(),
            "max_feasible_step: lower shape mismatch"
        );
        assert_eq!(
            self.shape(),
            upper.shape(),
            "max_feasible_step: upper shape mismatch"
        );
        let mut tau = f64::INFINITY;
        ndarray::Zip::from(self)
            .and(step)
            .and(lower)
            .and(upper)
            .for_each(|&x, &s, &l, &u| {
                let t = max_feasible_step_component(x, s, l, u);
                if t < tau {
                    tau = t;
                }
            });
        tau
    }

    fn cl_kkt_inf_norm(&self, d_sq: &Self) -> f64 {
        assert_eq!(
            self.shape(),
            d_sq.shape(),
            "cl_kkt_inf_norm: shape mismatch"
        );
        let mut best = 0.0_f64;
        ndarray::Zip::from(self).and(d_sq).for_each(|&v, &d| {
            let candidate = v.abs() / d;
            if candidate > best {
                best = candidate;
            }
        });
        best
    }

    fn weighted_norm_squared(&self, weights: &Self) -> f64 {
        assert_eq!(
            self.shape(),
            weights.shape(),
            "weighted_norm_squared: shape mismatch"
        );
        let mut sum = 0.0_f64;
        ndarray::Zip::from(self)
            .and(weights)
            .for_each(|&v, &w| sum += v * v * w);
        sum
    }

    fn project_strictly_inside(&mut self, lower: &Self, upper: &Self, rstep: f64) {
        assert_eq!(
            self.shape(),
            lower.shape(),
            "project_strictly_inside: lower shape mismatch"
        );
        assert_eq!(
            self.shape(),
            upper.shape(),
            "project_strictly_inside: upper shape mismatch"
        );
        ndarray::Zip::from(self)
            .and(lower)
            .and(upper)
            .for_each(|x, &l, &u| {
                *x = project_strictly_inside_component(*x, l, u, rstep);
            });
    }
}