//! Data analysis routines

use std::array;
use crate::{approx, algebra, num};
use crate::traits::*;
use crate::types::*;

/// Compute the covariance matrix of a list of 2D points
pub fn covariance_2 <S> (points : &[Point2 <S>]) -> Matrix2 <S> where
  S : Field + num::NumCast + MaybeSerDes
{
  let npoints      = S::from (points.len()).unwrap();
  let mean         = points.iter().map (|p| p.0).sum::<Vector2 <_>>() / npoints;
  let centered     = points.iter().map (|p| p.0 - mean);
  let [xx, xy, yy] = centered.map (|p| [p.x * p.x, p.x * p.y, p.y * p.y ])
    .fold ([S::zero(); 3], |acc, item| array::from_fn (|i| acc[i] + item[i]));
  let div = S::one() / (npoints - S::one());
  let cov_xx = xx * div;
  let cov_xy = xy * div;
  let cov_yx = cov_xy;
  let cov_yy = yy * div;
  Matrix2::from_col_arrays ([
    [cov_xx, cov_xy],
    [cov_yx, cov_yy]
  ])
}

/// Compute the covariance matrix of a list of 3D points
pub fn covariance_3 <S> (points : &[Point3 <S>]) -> Matrix3 <S> where
  S : Field + num::NumCast + MaybeSerDes
{
  let npoints      = S::from (points.len()).unwrap();
  let mean         = points.iter().map (|p| p.0).sum::<Vector3 <_>>() / npoints;
  let centered     = points.iter().map (|p| p.0 - mean);
  let [xx, xy, xz, yy, yz, zz] = centered.map (|p| [
    p.x * p.x,
    p.x * p.y,
    p.x * p.z,
    p.y * p.y,
    p.y * p.z,
    p.z * p.z
  ]).fold ([S::zero(); 6], |acc, item| array::from_fn (|i| acc[i] + item[i]));
  let div = S::one() / (npoints - S::one());
  Matrix3::from_col_arrays ([
    [xx, xy, xz],
    [xy, yy, yz],
    [xz, yz, zz]
  ]) * div
}

/// 2D principal component analysis
pub fn pca_2 <S> (points : &[Point2 <S>]) -> Rotation2 <S> where
  S : Real + num::NumCast + approx::RelativeEq <Epsilon=S> + std::fmt::Debug
    + MaybeSerDes,
  Vector2 <S> : NormedVectorSpace <S>
{
  let covariance = covariance_2 (points);
  // covariance matrix is symmetric, therefore eigenvectors should be orthogonal and
  // thus point in the direction of the principal components
  let (eigenvalues, mut eigenvectors) : (Vec<_>, Vec<_>) =
    algebra::linear::eigensystem_2 (covariance).iter().copied().map (Option::unwrap)
      .unzip();
  let eigenvectors = {
    // normalize eigenvectors and create array
    eigenvectors.iter_mut().for_each (|v| *v = NormedVectorSpace::normalize (*v).into());
    <[Vector2 <S>; 2]>::try_from (eigenvectors.as_slice()).unwrap()
      .map (Vector2::into_array)
  };
  // the principal component with the most variance corresponds to the largest
  // eigenvalue, we take this to be the first principal component
  debug_assert!(eigenvalues[0] >= eigenvalues[1]);
  let mut mat = Matrix2::from_col_arrays (eigenvectors);
  // ensure determinant == 1
  if mat.determinant().is_negative() {
    mat.cols[0] *= -S::one();
  }
  Rotation2::new_approx (mat).unwrap()
}

/// 3D principal component analysis
pub fn pca_3 <S> (points : &[Point3 <S>]) -> Rotation3 <S> where
  S : Real + num::NumCast + approx::RelativeEq <Epsilon=S> + std::fmt::Debug
    + MaybeSerDes,
  Vector2 <S> : NormedVectorSpace <S>
{
  let covariance = covariance_3 (points);
  // covariance matrix is symmetric, therefore eigenvectors should be orthogonal and
  // thus point in the direction of the principal components
  let (eigenvalues, mut eigenvectors) : (Vec<_>, Vec<_>) =
    algebra::linear::eigensystem_3 (covariance).iter().flatten().copied().unzip();
  let eigenvectors = {
    // normalize eigenvectors and create array
    eigenvectors.iter_mut().for_each (|v| *v = NormedVectorSpace::normalize (*v).into());
    <[Vector3 <S>; 3]>::try_from (eigenvectors.as_slice()).unwrap()
      .map (Vector3::into_array)
  };
  // the principal component with the most variance corresponds to the largest
  // eigenvalue, we take this to be the first principal component
  debug_assert!(eigenvalues[0] >= eigenvalues[1]);
  debug_assert!(eigenvalues[1] >= eigenvalues[2]);
  let mut mat = Matrix3::from_col_arrays (eigenvectors);
  // ensure determinant == 1
  if mat.determinant().is_negative() {
    mat.cols[0] *= -S::one();
  }
  Rotation3::new_approx (mat).unwrap()
}

#[cfg(test)]
mod tests {
  #![expect(clippy::unreadable_literal)]

  use super::*;
  #[test]
  fn covariance_2d() {
    let points = [
      [ 1.0, 2.0],
      [ 1.0, 3.0],
      [10.0, 2.0]
    ].map (Point2::<f32>::from);
    let covariance = covariance_2 (points.as_slice());
    assert_eq!(covariance.into_col_arrays(),
      [ [27.0, -1.5],
        [-1.5, 0.3333333] ]);
  }

  #[test]
  fn pca_2d() {
    let points = [
      [-1.0, -1.0],
      [-1.0,  1.0],
      [ 1.0, -1.0],
      [ 1.0,  1.0]
    ].map (Point2::<f32>::from);
    let components = pca_2 (points.as_slice()).into_col_arrays();
    assert_eq!(components[0], [1.0, 0.0]);
    assert_eq!(components[1], [0.0, 1.0]);
  }

  #[test]
  fn pca_3d() {
    use std::f32::consts::{FRAC_1_SQRT_2, FRAC_PI_4};
    let points = [
      [ 1.0, -1.0, -1.0],
      [ 1.0, -1.0,  1.0],
      [ 1.0,  1.0, -1.0],
      [ 1.0,  1.0,  1.0],
      [-1.0, -1.0, -1.0],
      [-1.0, -1.0,  1.0],
      [-1.0,  1.0, -1.0],
      [-1.0,  1.0,  1.0]
    ].map (Point3::<f32>::from);
    let components = pca_3 (points.as_slice()).into_col_arrays();
    assert_eq!(components[0], [1.0, 0.0, 0.0]);
    assert_eq!(components[1], [0.0, 1.0, 0.0]);
    assert_eq!(components[2], [0.0, 0.0, 1.0]);

    // this case returns axis-aligned components, this agrees when numpy output (except
    // that numpy returns the identity matrix)
    let rotation = Rotation3::from_angle_y (FRAC_PI_4.into());
    let points = [
      [ 1.0, -1.0, -1.0],
      [ 1.0, -1.0,  1.0],
      [ 1.0,  1.0, -1.0],
      [ 1.0,  1.0,  1.0],
      [-1.0, -1.0, -1.0],
      [-1.0, -1.0,  1.0],
      [-1.0,  1.0, -1.0],
      [-1.0,  1.0,  1.0]
    ].map (Point3::<f32>::from).map (|p| rotation.rotate (p));
    let components = pca_3 (points.as_slice()).into_col_arrays();
    assert_eq!(components[0], [0.0, -1.0, 0.0]);
    assert_eq!(components[1], [1.0,  0.0, 0.0]);
    assert_eq!(components[2], [0.0,  0.0, 1.0]);

    let rotation = Rotation3::from_angle_y (FRAC_PI_4.into());
    let points = [
      [ 1.0, -1.0, -4.0],
      [ 1.0, -1.0,  4.0],
      [ 1.0,  1.0, -4.0],
      [ 1.0,  1.0,  4.0],
      [-1.0, -1.0, -4.0],
      [-1.0, -1.0,  4.0],
      [-1.0,  1.0, -4.0],
      [-1.0,  1.0,  4.0]
    ].map (Point3::<f32>::from).map (|p| rotation.rotate (p));
    let components = pca_3 (points.as_slice());
    approx::assert_relative_eq!(components.mat(), Matrix3::from_col_arrays ([
      [FRAC_1_SQRT_2, 0.0, FRAC_1_SQRT_2],
      [0.0, 1.0, 0.0],
      [-FRAC_1_SQRT_2, 0.0, FRAC_1_SQRT_2]
    ]));
  }
}
