//! Perform PCA on a collection of vectors.
//! Assumes the full set of data can be loaded into memory.
use std::f32;
use linxal::prelude::*;
use ndarray::prelude::*;
use ndarray::{Data, Ix2};
use model::{Matrix, Vector};
#[derive(Debug)]
pub enum PCAError {
    BadSVD(SVDError),
    BadTarget
}

pub enum PCATarget {
    Dimension(usize),
    ExplainedVariance(f32)
}

/// Return true iff the target is valid, given the dimension of the
/// input data.
fn validate_target(target: &PCATarget, data_dim: usize) -> bool {
    match *target {
        PCATarget::Dimension(n) =>  n <= data_dim ,
        PCATarget::ExplainedVariance(p) =>  (0.0 < p) && (p <= 1.0)
    }
}

/// Returned the number of dimensions needed to satisfy the PCA
/// target, given the eigenvectors.
fn num_eigenvectors(target: &PCATarget, eig: &[f32]) -> usize {
    match target {
        &PCATarget::Dimension(n) => n,
        &PCATarget::ExplainedVariance(p) => {
            let total_var = eig.iter().map(|v: &f32| *v * *v).sum::<f32>();
            let mut explained_var = 0.0;
            let n = eig.len();
            for i in 0..n {
                explained_var += eig[i] * eig[i];
                if explained_var / total_var >= p {
                    return i+1
                }
            }
            n
        }
    }
}

/// Mean adjust a matrix by subtracting off the mean from every row,
/// returning a new matrix.
fn mean_adjusted<D1, D2>(m: &ArrayBase<D1, Ix2>, mean: &ArrayBase<D2, Ix>) -> Matrix<f32>
    where D1: Data<Elem=f32>, D2: Data<Elem=f32> {
    let copy = m.to_owned();
    copy - mean
}

// fn into_mean_adjusted(mut m: Matrix<f32>, mean: &Vector<f32>) -> Matrix<f32> {
//     for mut row in m.outer_iter_mut() {
//         for (i, e) in row.iter_mut().enumerate() {
//             *e -= mean[i];
//         }
//     }
//     m
// }


pub struct PCA {
    eigenvectors: Matrix<f32>,
    mean: Vector<f32>,
}

enum PCAMethod {
    SVD,
    CovarSVD
}

impl PCA {
    pub fn new(data: &Matrix<f32>, target: PCATarget) -> Result<PCA, PCAError> {
        if !validate_target(&target, data.cols()) {
            return Err(PCAError::BadTarget);
        }

        // Choose the method based on the size of the dataset.
        let method = if data.rows() <= 200000 {
            PCAMethod::SVD
        } else {
            PCAMethod::CovarSVD
        };

        // Compute the mean of the data.
        let mean: Vector<f32> = data.mean(Axis(0));

        // Create a new data matrix by subtracting the mean from each
        // row.
        let mean_adjusted_data = mean_adjusted(data, &mean);

        let a = match method {
            PCAMethod::CovarSVD => {
                // Compute the covariance matrix;
                let d = data.cols();
                let mut covar: Matrix<f32> = Matrix::zeros((d, d));
                for row in mean_adjusted_data.outer_iter() {
                    for i in 0..d {
                        for j in 0.. d {
                            covar[[i, j]] += row[i] * row[j];
                        }
                    }
                }

                println!("Covar: {:?}", covar);

                // Compute the eigenvalues and eigenvectors.
                let val2 = SymEigen::compute_mut(&mut covar, Symmetric::Upper, true).ok().unwrap();
                let eigenvalues = val2.iter().map(|v| { v.sqrt() }).collect::<Vec<f32>>();

                (eigenvalues, covar)
            },
            PCAMethod::SVD => {
                println!("Taking svd of matrix, size {:?}", mean_adjusted_data.dim());

                match SVD::compute_into(mean_adjusted_data, false, true) {
                    Ok(solution) => (solution.values.iter().cloned().collect(), solution.right_vectors.unwrap()),
                    Err(e) => return Err(PCAError::BadSVD(e))
                }
            }
        };

        let (eigenvalues, mut vec) = a;

        let k = num_eigenvectors(&target, &eigenvalues);
        let ev_subset = vec.view_mut().reversed_axes().split_at(Axis(1), k).0.to_owned();

        Ok( PCA { eigenvectors: ev_subset, mean: mean })

    }

    /// Transform a set of data into its reduced dimensionas
    pub fn transform_data<D>(&self, data: &ArrayBase<D, Ix2>) -> Matrix<f32> where D: Data<Elem=f32> {
        mean_adjusted(data, &self.mean).dot(&self.eigenvectors)
    }

    /// Transform an individual vector.
    pub fn transform_datum<D>(&self, datum: &ArrayBase<D, Ix>) -> Vector<f32> where D: Data<Elem=f32> {
        let v: Vector<f32> = datum - &self.mean;
        self.eigenvectors.t().dot(&v)
    }

    /// Reconstruct the best approximation of the original vector using
    /// the transformed vector.
    pub fn reconstruct_datum(&self, datum: &Vector<f32>) -> Vector<f32> {
        self.eigenvectors.dot(datum) + &self.mean
    }

    /// Reconstruct the best approximatio of the original vector using
    /// a subset of the elements.
    pub fn reconstruct_datum_partial(&self, datum: &Vector<f32>, dim: usize) -> Vector<f32> {
        assert!(dim <= self.dim());
        println!("EV; {:?}", self.eigenvectors.dim());
        let m = self.eigenvectors.view().split_at(Axis(1), dim).0;
        println!("EV sub: {:?}", m.dim());
        let v = datum.view().split_at(Axis(0), dim).0;
        println!("V sub: {:?}", v.dim());
        // let mu = self.mean.view().split_at(Axis(0), dim).0;
        // println!("μ sub: {:?}", v.dim());

        m.dot(&v) + &self.mean
    }


    /// Reconstruct the best approximation for a data set using
    /// the transformed data set.
    // pub fn reconstruct_data(&self, data: &Matrix<f32>) -> Matrix<f32> {
    //     let r = &self.eigenvectors * data.transpose();
    //     let nm = -&self.mean;
    //     into_mean_adjusted(r, &nm)
    // }

    /// Return the matrix of eigenvectors, with each column
    /// representing an eigenvector.
    pub fn eigenvectors(&self) -> &Matrix<f32> {
        &self.eigenvectors
    }

    pub fn dim(&self) -> usize {
        self.eigenvectors.cols()
    }
}

#[test]
fn variance_eig_test()
{
    let e: Vec<f32> = [0.4, 0.3, 0.2, 0.1].iter().map(|x| (*x as f32).sqrt()).collect();
    let eig: Vector<f32> = Vector::new(e);
    assert_eq!(num_eigenvectors(&PCATarget::ExplainedVariance(0.5), &eig.data()), 2);
    assert_eq!(num_eigenvectors(&PCATarget::ExplainedVariance(0.8), &eig.data()), 3);
    assert_eq!(num_eigenvectors(&PCATarget::ExplainedVariance(0.85), &eig.data()), 3);
    assert_eq!(num_eigenvectors(&PCATarget::ExplainedVariance(0.91), &eig.data()), 4);
}