//! Helper functions for Principal Component type analyses with implementations
//! for randomised SVD on dense or sparse matrices.

use faer::{Mat, MatMut, MatRef};
use num_traits::Float;
use rand::prelude::*;
use rand_distr::Normal;
use rayon::prelude::*;

use super::*;
use crate::prelude::*;

////////////////
// Structures //
////////////////

/// Structure for random SVD results
#[derive(Clone, Debug)]
pub struct RandomSvdResults<T> {
    /// Matrix u of the SVD decomposition
    pub u: Mat<T>,
    /// Matrix v of the SVD decomposition
    pub v: Mat<T>,
    /// Eigen vectors of the SVD decomposition
    pub s: Vec<T>,
}

/// Structure for SVD results
///
/// ### Fields
///
/// * `u` - Matrix u of the SVD decomposition
/// * `v` - Matrix v of the SVD decomposition
/// * `s` - Eigen vectors of the SVD decomposition
#[derive(Clone, Debug)]
pub struct SvdResults<T> {
    /// Matrix u of the SVD decomposition
    pub u: Mat<T>,
    /// Matrix v of the SVD decomposition
    pub v: Mat<T>,
    /// Eigen vectors of the SVD decomposition
    pub s: Vec<T>,
}

/// Trait to return the different matrices from the Svd Resuls
pub trait SvdResult<T> {
    /// Returns the matrix u of the SVD decomposition
    fn u(&self) -> &faer::Mat<T>;
    /// Returns the matrix v of the SVD decomposition
    fn v(&self) -> &faer::Mat<T>;
    /// Returns the eigen vectors of the SVD decomposition
    fn s(&self) -> &[T];
}

/// Implementations of the SvdResult trait for randomised SvdResults
impl<T> SvdResult<T> for RandomSvdResults<T> {
    fn u(&self) -> &faer::Mat<T> {
        &self.u
    }
    fn v(&self) -> &faer::Mat<T> {
        &self.v
    }
    fn s(&self) -> &[T] {
        &self.s
    }
}

/// Implementations of the SvdResult trait for SvdResults
impl<T> SvdResult<T> for SvdResults<T> {
    fn u(&self) -> &faer::Mat<T> {
        &self.u
    }
    fn v(&self) -> &faer::Mat<T> {
        &self.v
    }
    fn s(&self) -> &[T] {
        &self.s
    }
}

/////////////
// Helpers //
/////////////

/// Calculate the principal component scores from the SVD results
///
/// ### Params
///
/// * `svd_results` - The (randomised) SVD results
///
/// ### Returns
///
/// The principal component scores
pub fn compute_pc_scores<T, S>(svd_results: &S) -> Mat<T>
where
    T: Float,
    S: SvdResult<T>,
{
    let n_cells = svd_results.u().nrows();
    let n_pcs = svd_results.s().len();

    Mat::from_fn(n_cells, n_pcs, |i, j| {
        svd_results.u()[(i, j)] * svd_results.s()[j]
    })
}

///////////////
// Functions //
///////////////

/// Get the eigenvalues and vectors from a covar or cor matrix
///
/// Function will panic if the matrix is not symmetric
///
/// ### Params
///
/// * `matrix` - The correlation or co-variance matrix
/// * `top_n` - How many of the top eigen vectors and values to return.
///
/// ### Returns
///
/// A vector of tuples corresponding to the top eigen pairs.
pub fn get_top_eigenvalues<T>(
    matrix: &Mat<T>,
    top_n: usize,
) -> Result<Vec<(T, Vec<T>)>, BixverseErrors>
where
    T: BixverseFloat,
{
    // Ensure the matrix is square
    assert_symmetric_mat!(matrix);

    let eigendecomp = matrix.eigen().map_err(|_| BixverseErrors::FaerEigenError)?;

    let s = eigendecomp.S();
    let u = eigendecomp.U();

    // Extract the real part of the eigenvalues and vectors
    let mut eigenpairs = s
        .column_vector()
        .iter()
        .zip(u.col_iter())
        .map(|(l, v)| {
            let l_real = l.re;
            let v_real = v.iter().map(|v_i| v_i.re).collect::<Vec<T>>();
            (l_real, v_real)
        })
        .collect::<Vec<(T, Vec<T>)>>();

    // Sort and return Top N
    eigenpairs.sort_by(|a, b| b.0.total_cmp(&a.0));

    let res: Vec<(T, Vec<T>)> = eigenpairs.into_iter().take(top_n).collect();

    Ok(res)
}

/// Randomised SVD
///
/// ### Params
///
/// * `x` - The matrix on which to apply the randomised SVD.
/// * `rank` - The target rank of the approximation (number of singular values,
///   vectors to compute).
/// * `seed` - Random seed for reproducible results.
/// * `oversampling` - Additional samples beyond the target rank to improve
///   accuracy. Defaults to 10 if not specified.
/// * `n_power_iter` - Number of power iterations to perform for better
///   approximation quality. More iterations generally improve accuracy but
///   increase computation time. Defaults to 2 if not specified.
///
/// ### Returns
///
/// The randomised SVD results in form of `RandomSvdResults`.
///
/// ### Algorithm Details
///
/// 1. Generate a random Gaussian matrix Ω of size n × (rank + oversampling)
/// 2. Compute Y = X * Ω to capture the range of X
/// 3. Orthogonalize Y using QR decomposition to get Q
/// 4. Apply power iterations: for each iteration, compute Z = X^T * Q, then Q = QR(X * Z)
/// 5. Form B = Q^T * X and compute its SVD
/// 6. Reconstruct the final SVD: U = Q * U_B, V = V_B, S = S_B
pub fn randomised_svd<T>(
    x: MatRef<T>,
    rank: usize,
    seed: usize,
    oversampling: Option<usize>,
    n_power_iter: Option<usize>,
) -> Result<RandomSvdResults<T>, BixverseErrors>
where
    T: BixverseFloat,
{
    let ncol = x.ncols();
    let nrow = x.nrows();
    let os = oversampling.unwrap_or(DEFAULT_OVERSAMPLING_RAND_SVD);
    let sample_size = (rank + os).min(ncol.min(nrow));
    let n_iter = n_power_iter.unwrap_or(DEFAULT_N_POWER_ITERS_RAND_SVD);

    let mut rng = StdRng::seed_from_u64(seed as u64);
    let normal = Normal::new(0.0, 1.0).unwrap();
    let omega = Mat::from_fn(ncol, sample_size, |_, _| {
        T::from_f64(normal.sample(&mut rng)).unwrap()
    });

    let y = x * omega;
    let mut q = y.qr().compute_thin_Q();
    for _ in 0..n_iter {
        let z = x.transpose() * q;
        q = (x * z).qr().compute_thin_Q();
    }

    let b = q.transpose() * x;
    let svd = b
        .thin_svd()
        .map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;

    Ok(RandomSvdResults {
        u: q * svd.U(),
        v: svd.V().cloned(),
        s: svd.S().column_vector().iter().copied().collect(),
    })
}

///////////////////////////
// Sparse randomised SVD //
///////////////////////////

/// Randomised sparse SVD - never forms dense intermediate matrices
///
/// ### Params
///
/// * `matrix` - Sparse matrix (CSR or CSC), consumed
/// * `rank` - Target rank
/// * `seed` - For reproducibility
/// * `use_second_layer` - Whether to use the second layer of the sparse matrix
///   for SVD calculation.
/// * `oversampling` - Additional samples (default 10)
/// * `n_power_iter` - Power iterations for accuracy (default 2)
/// * `col_means` - Optional column means for implicit mean centering
/// * `col_stds` - Optional column sds for implicit variance normalising
/// * `row_offsets` - Additional offsets (for example for CLR-type PCA in single
///   cell).
///
/// ### Returns
///
/// `RandomSvdResults` containing U (n x k), S (length k), and V (m x k)
#[allow(clippy::too_many_arguments)]
pub fn randomised_sparse_svd<T, F>(
    matrix: CompressedSparseData2<T>,
    rank: usize,
    seed: u64,
    use_second_layer: bool,
    oversampling: Option<usize>,
    n_power_iter: Option<usize>,
    col_means: Option<&[F]>,
    col_stds: Option<&[F]>,
    row_offsets: Option<&[F]>,
) -> Result<RandomSvdResults<F>, BixverseErrors>
where
    T: BixverseNumeric + Into<F>,
    F: BixverseFloat,
{
    let (n, m) = matrix.shape;
    let os = oversampling.unwrap_or(DEFAULT_OVERSAMPLING_RAND_SVD);
    let sample_size = (rank + os).min(m).min(n);
    let n_iter = n_power_iter.unwrap_or(DEFAULT_N_POWER_ITERS_RAND_SVD);

    // consume the input: transform if needed, drop the original CSC.
    let csr = match matrix.cs_type {
        CompressedSparseFormat::Csr => matrix,
        CompressedSparseFormat::Csc => matrix.transform_single_layer(use_second_layer)?,
    };

    // pick the active data slice without materialising an f64 copy.
    // values are converted to F at access sites (free cast for f32 -> f64).
    let active_data: &[T] = if use_second_layer {
        csr.data_2
            .as_ref()
            .ok_or(BixverseErrors::Data2NotAvailable)?
            .as_slice()
    } else {
        csr.data.as_slice()
    };

    // helper: read a non-zero value as f64 for stable accumulation.
    let val_f64 = |idx: usize| -> f64 {
        let v: F = active_data[idx].into();
        v.to_f64().unwrap()
    };

    // pre-divide input (m x ncols) by col_stds once, avoiding per-nonzero division.
    let prescale = |x: MatRef<F>| -> Option<Mat<F>> {
        col_stds.map(|sd| Mat::from_fn(x.nrows(), x.ncols(), |i, col| x[(i, col)] / sd[i]))
    };

    // y = A * x_scaled - 1 * (mu^T * x_scaled)
    // expects x_scaled to already be divided by sigma if applicable.
    // accumulates per-row in f64 for numerical stability.
    let sparse_matvec_a = |x_scaled: MatRef<F>, y: MatMut<F>| {
        let ncols = x_scaled.ncols();

        let mean_dots: Vec<f64> = if let Some(mu) = col_means {
            (0..ncols)
                .map(|col| {
                    (0..m)
                        .map(|j| mu[j].to_f64().unwrap() * x_scaled[(j, col)].to_f64().unwrap())
                        .sum::<f64>()
                })
                .collect()
        } else {
            Vec::new()
        };

        let x_sums: Vec<f64> = if row_offsets.is_some() {
            (0..ncols)
                .map(|col| {
                    (0..m)
                        .map(|j| x_scaled[(j, col)].to_f64().unwrap())
                        .sum::<f64>()
                })
                .collect()
        } else {
            Vec::new()
        };

        let y_ptr = y.as_ptr_mut() as usize;
        let y_row_stride = y.row_stride();
        let y_col_stride = y.col_stride();

        (0..n).into_par_iter().for_each(|i| {
            let base = y_ptr as *mut F;
            let mut row_acc = vec![0f64; ncols];

            for idx in csr.indptr[i]..csr.indptr[i + 1] {
                let j = csr.indices[idx as usize] as usize;
                let idx_usize = idx as usize;
                let a_val = val_f64(idx_usize);
                for col in 0..ncols {
                    row_acc[col] += a_val * x_scaled[(j, col)].to_f64().unwrap();
                }
            }

            if col_means.is_some() {
                for col in 0..ncols {
                    row_acc[col] -= mean_dots[col];
                }
            }

            if let Some(off) = row_offsets {
                let m_i = off[i].to_f64().unwrap();
                for col in 0..ncols {
                    row_acc[col] -= m_i * x_sums[col];
                }
            }

            for col in 0..ncols {
                unsafe {
                    let ptr = base.offset(i as isize * y_row_stride + col as isize * y_col_stride);
                    *ptr = F::from_f64(row_acc[col]).unwrap();
                }
            }
        });
    };

    let sparse_matvec_at = |x: MatRef<F>, mut y: MatMut<F>| {
        let ncols = x.ncols();

        let m_dot_x: Vec<f64> = if let Some(off) = row_offsets {
            (0..ncols)
                .map(|col| {
                    (0..n)
                        .map(|i| off[i].to_f64().unwrap() * x[(i, col)].to_f64().unwrap())
                        .sum::<f64>()
                })
                .collect()
        } else {
            Vec::new()
        };

        let (result, col_sums) = (0..n)
            .into_par_iter()
            .fold(
                || (vec![0f64; m * ncols], vec![0f64; ncols]),
                |(mut acc, mut cs), i| {
                    for col in 0..ncols {
                        cs[col] += x[(i, col)].to_f64().unwrap();
                    }
                    for idx in csr.indptr[i]..csr.indptr[i + 1] {
                        let idx_usize = idx as usize;
                        let j = csr.indices[idx_usize] as usize;
                        let a_val = val_f64(idx_usize);
                        for col in 0..ncols {
                            acc[j * ncols + col] += a_val * x[(i, col)].to_f64().unwrap();
                        }
                    }
                    (acc, cs)
                },
            )
            .reduce(
                || (vec![0f64; m * ncols], vec![0f64; ncols]),
                |(mut a, mut cs_a), (b, cs_b)| {
                    for i in 0..a.len() {
                        a[i] += b[i];
                    }
                    for i in 0..ncols {
                        cs_a[i] += cs_b[i];
                    }
                    (a, cs_a)
                },
            );

        for j in 0..m {
            for col in 0..ncols {
                let mut val = result[j * ncols + col];
                if let Some(mu) = col_means {
                    val -= mu[j].to_f64().unwrap() * col_sums[col];
                }
                if row_offsets.is_some() {
                    val -= m_dot_x[col];
                }
                if let Some(sd) = col_stds {
                    val /= sd[j].to_f64().unwrap();
                }
                y[(j, col)] = F::from_f64(val).unwrap();
            }
        }
    };

    let mut rng = StdRng::seed_from_u64(seed);
    let normal = Normal::new(0.0, 1.0).unwrap();
    let omega = Mat::from_fn(m, sample_size, |_, _| {
        F::from_f64(normal.sample(&mut rng)).unwrap()
    });

    let omega_scaled = prescale(omega.as_ref());
    let mut y = Mat::<F>::zeros(n, sample_size);
    sparse_matvec_a(
        omega_scaled
            .as_ref()
            .map(|m| m.as_ref())
            .unwrap_or(omega.as_ref()),
        y.as_mut(),
    );
    drop(omega);
    drop(omega_scaled);

    let mut q = y.qr().compute_thin_Q();

    let mut z = Mat::<F>::zeros(m, sample_size);
    let mut y_new = Mat::<F>::zeros(n, sample_size);

    for _ in 0..n_iter {
        sparse_matvec_at(q.as_ref(), z.as_mut());
        let z_scaled = prescale(z.as_ref());
        sparse_matvec_a(
            z_scaled.as_ref().map(|m| m.as_ref()).unwrap_or(z.as_ref()),
            y_new.as_mut(),
        );
        q = y_new.qr().compute_thin_Q();
    }

    // memory dropping
    drop(z);
    drop(y_new);
    drop(y);

    let mut b_t = Mat::<F>::zeros(m, sample_size);
    sparse_matvec_at(q.as_ref(), b_t.as_mut());
    let b = b_t.transpose().to_owned();

    let svd = b
        .thin_svd()
        .map_err(|e| BixverseErrors::FaerSvdError(format!("{e:?}")))?;

    let u = &q * svd.U();
    let s: Vec<F> = svd.S().column_vector().iter().copied().collect();
    let v = svd.V().to_owned();

    Ok(RandomSvdResults { u, s, v })
}

///////////
// Tests //
///////////

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

    #[test]
    fn test_compute_pc_scores() {
        let u: Mat<f64> = Mat::from_fn(2, 2, |i, j| if i == j { 1.0 } else { 0.0 });
        let v: Mat<f64> = Mat::from_fn(2, 2, |i, j| if i == j { 1.0 } else { 0.0 });
        let s = vec![3.0, 1.5];

        let svd_res = SvdResults { u, v, s };
        let scores = compute_pc_scores(&svd_res);

        // PC scores = U * S
        assert!((scores[(0, 0)] - 3.0).abs() < 1e-6);
        assert!((scores[(0, 1)] - 0.0).abs() < 1e-6);
        assert!((scores[(1, 0)] - 0.0).abs() < 1e-6);
        assert!((scores[(1, 1)] - 1.5).abs() < 1e-6);
    }

    #[test]
    fn test_get_top_eigenvalues() {
        let mat: Mat<f64> = Mat::from_fn(
            2,
            2,
            |i, j| if i == j { 3.0 - (i as f64 * 2.0) } else { 0.0 },
        );
        let top_eigen = get_top_eigenvalues(&mat, 2).unwrap();

        assert_eq!(top_eigen.len(), 2);
        assert!((top_eigen[0].0 - 3.0).abs() < 1e-6);
        assert!((top_eigen[1].0 - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_randomised_svd_logic() {
        // Create a dense rank-1 matrix A = x * y^T
        // x = [1.0, 2.0, 3.0, 4.0]^T
        // y = [1.0, 0.5, 0.25]^T
        let x = [1.0, 2.0, 3.0, 4.0];
        let y = [1.0, 0.5, 0.25];
        let mut mat: Mat<f64> = Mat::zeros(4, 3);
        for i in 0..4 {
            for j in 0..3 {
                mat[(i, j)] = x[i] * y[j];
            }
        }

        // We only need the top 1 PC
        let svd = randomised_svd(mat.as_ref(), 1, 42, Some(5), Some(4)).unwrap();

        // Test U (Left Singular Vector) correlation with x
        let u_col = svd.u.col(0);
        let x_norm: f64 = x.iter().map(|v| v * v).sum::<f64>().sqrt();
        let mut dot_u = 0.0;
        for i in 0..4 {
            dot_u += u_col[i] * (x[i] / x_norm);
        }

        // Test V (Right Singular Vector) correlation with y
        let v_col = svd.v.col(0);
        let y_norm: f64 = y.iter().map(|v| v * v).sum::<f64>().sqrt();
        let mut dot_v = 0.0;
        for j in 0..3 {
            dot_v += v_col[j] * (y[j] / y_norm);
        }

        // The absolute correlation should be > 0.999 (allowing for sign flips)
        assert!(dot_u.abs() > 0.999);
        assert!(dot_v.abs() > 0.999);
    }

    #[test]
    fn test_randomised_sparse_svd_logic() {
        // Create a sparse rank-1 matrix A = x * y^T
        // x = [0.0, 2.0, 0.0, 4.0]^T
        // y = [1.0, 0.0, 0.5]^T
        // Non-zeros will only exist where x_i != 0 AND y_j != 0
        let data = vec![2.0, 1.0, 4.0, 2.0];
        let indices = vec![0, 2, 0, 2];
        let indptr = vec![0, 0, 2, 2, 4];
        let shape = (4, 3);

        let csr = CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, shape);

        let no_params: Option<&[f64]> = None;
        let svd = randomised_sparse_svd(
            csr,
            1,
            42,
            false,
            Some(5),
            Some(4),
            no_params,
            no_params,
            None,
        )
        .unwrap();

        let u_col = svd.u.col(0);
        let x_norm = (2.0_f64.powi(2) + 4.0_f64.powi(2)).sqrt(); // sqrt(20)
        let dot_u = (u_col[1] * 2.0 + u_col[3] * 4.0) / x_norm;
        assert!(dot_u.abs() > 0.999);

        let v_col = svd.v.col(0);
        let y_norm = (1.0_f64.powi(2) + 0.5_f64.powi(2)).sqrt(); // sqrt(1.25)
        let dot_v = (v_col[0] * 1.0 + v_col[2] * 0.5) / y_norm;
        assert!(dot_v.abs() > 0.999);
    }
}
