// GPL-3.0-or-later
// Copyright 2024--present rsx-rs developers

//! `pca` command: streaming PCA of the depth matrix.

use crate::markers_table::{MarkersTableStream, ParserConfig};
use crate::source::MarkerStream;
use std::io::Write;
use std::path::Path;

pub struct PcaParams {
    pub markers_table_path: String,
    pub output_dir: String,
    pub min_depth: u16,
    pub n_components: Option<usize>,
}

/// Internal result of the core PCA computation. Shared by the file-based
/// writer and the in-memory Arrow path.
struct ComputedPca {
    eigenvalues: Vec<f64>,
    eigenvectors: Vec<f64>,
    sorted_indices: Vec<usize>,
    n_markers: u64,
    n_individuals: usize,
    n_components: usize,
    total_variance: f64,
    individual_names: Vec<String>,
}

/// Core streaming PCA against the on-disk TSV.
fn compute_pca(
    markers_table_path: &str,
    min_depth: u16,
    n_components: Option<usize>,
) -> Result<ComputedPca, Box<dyn std::error::Error>> {
    let table_path = Path::new(markers_table_path);
    let config = ParserConfig {
        store_sequence: false,
        store_depths: true,
        compute_groups: false,
        min_depth,
    };
    let stream = MarkersTableStream::open(table_path, None, config)?;
    compute_pca_with_source(&stream, n_components)
}

/// Core streaming PCA against any `MarkerStream`.
///
/// Builds the centered Gram as a **streaming contraction** \(G = X_c^\mathrm{T} X_c\):
/// one pass accumulates uncentered \(X^\mathrm{T} X\) (upper triangle) and column
/// means, then applies \(G \leftarrow X^\mathrm{T} X - m\mu\mu^\mathrm{T}\).
fn compute_pca_with_source<S: MarkerStream>(
    source: &S,
    n_components: Option<usize>,
) -> Result<ComputedPca, Box<dyn std::error::Error>> {
    let n = source.header().n_individuals as usize;

    log::info!(
        "PCA: streaming {} individuals, building {}x{} Gram (streaming contraction)",
        n,
        n,
        n
    );

    let (mut gram, mut mean, n_markers) = accumulate_gram_streaming(source, n)?;

    if n_markers == 0 {
        return Err("No markers found".into());
    }

    log::info!("PCA: {} markers streamed, centering Gram matrix", n_markers);

    for i in 0..n {
        for j in 0..i {
            gram[i * n + j] = gram[j * n + i];
        }
    }

    let nm = n_markers as f64;
    for m in &mut mean {
        *m /= nm;
    }
    for i in 0..n {
        for j in 0..n {
            gram[i * n + j] -= nm * mean[i] * mean[j];
        }
    }

    log::info!("PCA: eigendecomposing {}x{} matrix", n, n);
    let (eigenvalues, eigenvectors) = jacobi_eigen(&mut gram, n);

    let mut indices: Vec<usize> = (0..n).collect();
    indices.sort_by(|&a, &b| eigenvalues[b].total_cmp(&eigenvalues[a]));

    let r = n_components.unwrap_or(n).min(n);
    let total_var: f64 = eigenvalues.iter().filter(|&&v| v > 0.0).sum();

    let header_cols = &source.header().columns;
    let individual_names: Vec<String> = (0..n)
        .map(|i| {
            if i + 2 < header_cols.len() {
                header_cols[i + 2].clone()
            } else {
                format!("ind{}", i + 1)
            }
        })
        .collect();

    Ok(ComputedPca {
        eigenvalues,
        eigenvectors,
        sorted_indices: indices,
        n_markers,
        n_individuals: n,
        n_components: r,
        total_variance: total_var,
        individual_names,
    })
}

/// Streaming contraction pass: uncentered Gram upper triangle + column sums.
type GramStreamingAcc = (Vec<f64>, Vec<f64>, u64);

fn accumulate_gram_streaming<S: MarkerStream>(
    source: &S,
    n: usize,
) -> Result<GramStreamingAcc, Box<dyn std::error::Error>> {
    let mut gram = vec![0.0f64; n * n];
    let mut mean = vec![0.0f64; n];
    let mut n_markers = 0u64;
    source.for_each(|marker| {
        if marker.n_individuals == 0 {
            return;
        }
        n_markers += 1;
        let depths = &marker.individual_depths;
        for i in 0..n {
            let xi = depths[i] as f64;
            mean[i] += xi;
            let base = i * n;
            for j in i..n {
                gram[base + j] += xi * (depths[j] as f64);
            }
        }
    })?;
    Ok((gram, mean, n_markers))
}

/// Sparse rank-1 outer product for highly zero depth rows (exact). Exposed for
/// microbenchmarks; production PCA uses the dense inlined loop above.
#[inline]
pub fn rank1_upper_update_sparse(gram: &mut [f64], mean: &mut [f64], depths: &[u16], n: usize) {
    debug_assert_eq!(depths.len(), n);
    for i in 0..n {
        let di = depths[i];
        if di == 0 {
            continue;
        }
        let xi = di as f64;
        mean[i] += xi;
        let base = i * n;
        for j in i..n {
            let dj = depths[j];
            if dj != 0 {
                gram[base + j] += xi * (dj as f64);
            }
        }
    }
}

#[cfg(test)]
mod gram_tests {
    use super::rank1_upper_update_sparse;

    #[test]
    fn sparse_rank1_matches_naive_and_beats_dense_on_sparse_rows() {
        let n = 64;
        let mut depths = vec![0u16; n];
        depths[1] = 3;
        depths[17] = 5;
        depths[40] = 2;
        depths[63] = 9;
        let mut g_s = vec![0.0; n * n];
        let mut m_s = vec![0.0; n];
        rank1_upper_update_sparse(&mut g_s, &mut m_s, &depths, n);
        let mut g_ref = vec![0.0; n * n];
        let mut m_ref = vec![0.0; n];
        for i in 0..n {
            let xi = depths[i] as f64;
            m_ref[i] += xi;
            for j in i..n {
                g_ref[i * n + j] += xi * (depths[j] as f64);
            }
        }
        for i in 0..n * n {
            assert!((g_s[i] - g_ref[i]).abs() < 1e-12);
        }
        let t0 = std::time::Instant::now();
        for _ in 0..30_000 {
            rank1_upper_update_sparse(&mut g_s, &mut m_s, &depths, n);
        }
        let sparse_ns = t0.elapsed().as_nanos();
        let mut g_d = vec![0.0; n * n];
        let mut m_d = vec![0.0; n];
        let t1 = std::time::Instant::now();
        for _ in 0..30_000 {
            for i in 0..n {
                let xi = depths[i] as f64;
                m_d[i] += xi;
                for j in i..n {
                    g_d[i * n + j] += xi * (depths[j] as f64);
                }
            }
        }
        let dense_ns = t1.elapsed().as_nanos();
        eprintln!("sparse_ns={sparse_ns} dense_ns={dense_ns}");
        assert!(sparse_ns < dense_ns, "sparse={sparse_ns} dense={dense_ns}");
    }
}

pub fn run(params: &PcaParams) -> Result<(), Box<dyn std::error::Error>> {
    let c = compute_pca(
        &params.markers_table_path,
        params.min_depth,
        params.n_components,
    )?;
    write_pca_outputs(&c, &params.output_dir)
}

pub fn run_with_source<S: MarkerStream>(
    source: &S,
    params: &PcaParams,
) -> Result<(), Box<dyn std::error::Error>> {
    let c = compute_pca_with_source(source, params.n_components)?;
    write_pca_outputs(&c, &params.output_dir)
}

fn write_pca_outputs(c: &ComputedPca, output_dir: &str) -> Result<(), Box<dyn std::error::Error>> {
    std::fs::create_dir_all(output_dir)?;

    let eigenval_path = Path::new(output_dir).join("eigenvalues.tsv");
    let mut f = std::fs::File::create(&eigenval_path)?;
    writeln!(f, "component\teigenvalue\tvariance_fraction\tcumulative")?;
    let mut cumulative = 0.0;
    for (k, &idx) in c.sorted_indices.iter().take(c.n_components).enumerate() {
        let ev = c.eigenvalues[idx].max(0.0);
        let frac = if c.total_variance > 0.0 {
            ev / c.total_variance
        } else {
            0.0
        };
        cumulative += frac;
        writeln!(f, "PC{}\t{:.6}\t{:.6}\t{:.6}", k + 1, ev, frac, cumulative)?;
    }

    let loadings_path = Path::new(output_dir).join("loadings.tsv");
    let mut f = std::fs::File::create(&loadings_path)?;
    write!(f, "individual")?;
    for k in 0..c.n_components {
        write!(f, "\tPC{}", k + 1)?;
    }
    writeln!(f)?;

    for (i, name) in c.individual_names.iter().enumerate() {
        write!(f, "{}", name)?;
        for &idx in c.sorted_indices.iter().take(c.n_components) {
            write!(f, "\t{:.6}", c.eigenvectors[i * c.n_individuals + idx])?;
        }
        writeln!(f)?;
    }

    let summary_path = Path::new(output_dir).join("summary.txt");
    let mut f = std::fs::File::create(&summary_path)?;
    writeln!(f, "Streaming PCA of depth matrix")?;
    writeln!(f, "Markers: {}", c.n_markers)?;
    writeln!(f, "Individuals: {}", c.n_individuals)?;
    writeln!(f, "Components: {}", c.n_components)?;
    writeln!(f, "Total variance: {:.2}", c.total_variance)?;
    writeln!(f)?;
    writeln!(f, "Top components:")?;
    cumulative = 0.0;
    for (k, &idx) in c
        .sorted_indices
        .iter()
        .take(c.n_components.min(10))
        .enumerate()
    {
        let ev = c.eigenvalues[idx].max(0.0);
        let frac = if c.total_variance > 0.0 {
            ev / c.total_variance
        } else {
            0.0
        };
        cumulative += frac;
        writeln!(
            f,
            "  PC{}: {:.4} variance ({:.1}% cumulative)",
            k + 1,
            frac,
            cumulative * 100.0
        )?;
    }

    log::info!(
        "PCA done: {} markers, {} individuals, {} components -> {}",
        c.n_markers,
        c.n_individuals,
        c.n_components,
        output_dir
    );

    Ok(())
}

/// Jacobi eigendecomposition for a symmetric matrix.
fn jacobi_eigen(a: &mut [f64], n: usize) -> (Vec<f64>, Vec<f64>) {
    let mut v = vec![0.0f64; n * n];
    for i in 0..n {
        v[i * n + i] = 1.0;
    }

    let max_sweeps = 100;
    let tol = 1e-15;

    for _sweep in 0..max_sweeps {
        let mut max_off = 0.0f64;
        for i in 0..n {
            for j in (i + 1)..n {
                max_off = max_off.max(a[i * n + j].abs());
            }
        }
        if max_off < tol {
            break;
        }

        for i in 0..n {
            for j in (i + 1)..n {
                let aij = a[i * n + j];
                if aij.abs() < tol * 0.01 {
                    continue;
                }

                let aii = a[i * n + i];
                let ajj = a[j * n + j];
                let diff = ajj - aii;

                let t = if diff.abs() < tol {
                    1.0
                } else {
                    let tau = diff / (2.0 * aij);
                    let sign = if tau >= 0.0 { 1.0 } else { -1.0 };
                    sign / (tau.abs() + (1.0 + tau * tau).sqrt())
                };

                let c = 1.0 / (1.0 + t * t).sqrt();
                let s = t * c;

                a[i * n + i] = aii - t * aij;
                a[j * n + j] = ajj + t * aij;
                a[i * n + j] = 0.0;
                a[j * n + i] = 0.0;

                for k in 0..n {
                    if k == i || k == j {
                        continue;
                    }
                    let aki = a[k * n + i];
                    let akj = a[k * n + j];
                    a[k * n + i] = c * aki - s * akj;
                    a[i * n + k] = a[k * n + i];
                    a[k * n + j] = s * aki + c * akj;
                    a[j * n + k] = a[k * n + j];
                }

                for k in 0..n {
                    let vki = v[k * n + i];
                    let vkj = v[k * n + j];
                    v[k * n + i] = c * vki - s * vkj;
                    v[k * n + j] = s * vki + c * vkj;
                }
            }
        }
    }

    let eigenvalues: Vec<f64> = (0..n).map(|i| a[i * n + i]).collect();
    (eigenvalues, v)
}

#[cfg(feature = "arrow-output")]
use arrow::array::builder::{Float64Builder, StringBuilder};
#[cfg(feature = "arrow-output")]
use arrow::datatypes::{DataType, Field, Schema};
#[cfg(feature = "arrow-output")]
use arrow::record_batch::RecordBatch;

#[cfg(feature = "arrow-output")]
pub struct PcaArrowResult {
    pub eigenvalues: RecordBatch,
    pub loadings: RecordBatch,
    pub n_markers: u64,
    pub n_individuals: usize,
    pub n_components: usize,
    pub total_variance: f64,
}

#[cfg(feature = "arrow-output")]
pub fn run_to_arrow(params: &PcaParams) -> Result<PcaArrowResult, Box<dyn std::error::Error>> {
    let c = compute_pca(
        &params.markers_table_path,
        params.min_depth,
        params.n_components,
    )?;
    pca_to_arrow_batches(c)
}

#[cfg(feature = "arrow-output")]
pub fn run_to_arrow_with_source<S: MarkerStream>(
    source: &S,
    params: &PcaParams,
) -> Result<PcaArrowResult, Box<dyn std::error::Error>> {
    let c = compute_pca_with_source(source, params.n_components)?;
    pca_to_arrow_batches(c)
}

#[cfg(feature = "arrow-output")]
fn pca_to_arrow_batches(c: ComputedPca) -> Result<PcaArrowResult, Box<dyn std::error::Error>> {
    let eigen_schema = Schema::new(vec![
        Field::new("component", DataType::Utf8, false),
        Field::new("eigenvalue", DataType::Float64, false),
        Field::new("variance_fraction", DataType::Float64, false),
        Field::new("cumulative", DataType::Float64, false),
    ]);

    let n_rows = c.n_components;
    let mut comp_b = StringBuilder::with_capacity(n_rows, n_rows * 4);
    let mut ev_b = Float64Builder::with_capacity(n_rows);
    let mut frac_b = Float64Builder::with_capacity(n_rows);
    let mut cum_b = Float64Builder::with_capacity(n_rows);

    let mut cumulative = 0.0;
    for (k, &idx) in c.sorted_indices.iter().take(c.n_components).enumerate() {
        let ev = c.eigenvalues[idx].max(0.0);
        let frac = if c.total_variance > 0.0 {
            ev / c.total_variance
        } else {
            0.0
        };
        cumulative += frac;

        comp_b.append_value(format!("PC{}", k + 1));
        ev_b.append_value(ev);
        frac_b.append_value(frac);
        cum_b.append_value(cumulative);
    }

    let eigen_batch = RecordBatch::try_new(
        std::sync::Arc::new(eigen_schema),
        vec![
            std::sync::Arc::new(comp_b.finish()),
            std::sync::Arc::new(ev_b.finish()),
            std::sync::Arc::new(frac_b.finish()),
            std::sync::Arc::new(cum_b.finish()),
        ],
    )?;

    let mut loading_fields = vec![Field::new("individual", DataType::Utf8, false)];
    for k in 0..c.n_components {
        loading_fields.push(Field::new(format!("PC{}", k + 1), DataType::Float64, false));
    }
    let loadings_schema = Schema::new(loading_fields);

    let n_ind = c.n_individuals;
    let n_pc = c.n_components;
    let mut ind_b = StringBuilder::with_capacity(n_ind, n_ind * 16);
    let mut pc_builders: Vec<Float64Builder> = (0..n_pc)
        .map(|_| Float64Builder::with_capacity(n_ind))
        .collect();

    for (i, name) in c.individual_names.iter().enumerate() {
        ind_b.append_value(name);
        for (k, &idx) in c.sorted_indices.iter().take(c.n_components).enumerate() {
            pc_builders[k].append_value(c.eigenvectors[i * c.n_individuals + idx]);
        }
    }

    let mut loading_columns: Vec<std::sync::Arc<dyn arrow::array::Array>> =
        vec![std::sync::Arc::new(ind_b.finish())];
    for mut b in pc_builders {
        loading_columns.push(std::sync::Arc::new(b.finish()));
    }

    let loadings_batch =
        RecordBatch::try_new(std::sync::Arc::new(loadings_schema), loading_columns)?;

    Ok(PcaArrowResult {
        eigenvalues: eigen_batch,
        loadings: loadings_batch,
        n_markers: c.n_markers,
        n_individuals: c.n_individuals,
        n_components: c.n_components,
        total_variance: c.total_variance,
    })
}

#[cfg(all(test, feature = "arrow-output"))]
mod tests {
    use super::*;

    fn make_pca_test_data(dir: &std::path::Path) -> std::path::PathBuf {
        let path = dir.join("markers.tsv");
        let mut f = std::fs::File::create(&path).unwrap();
        writeln!(f, "#Number of markers : 4").unwrap();
        writeln!(f, "id\tsequence\tind1\tind2\tind3\tind4\tind5").unwrap();
        writeln!(f, "0\tATCGATCG\t10\t5\t8\t12\t7").unwrap();
        writeln!(f, "1\tGGGGAAAA\t15\t20\t10\t0\t0").unwrap();
        writeln!(f, "2\tCCCCTTTT\t0\t0\t0\t25\t30").unwrap();
        writeln!(f, "3\tAAAATTTT\t5\t0\t3\t8\t6").unwrap();
        path
    }

    #[test]
    fn run_to_arrow_matches_file_based_pca() {
        let dir = std::env::temp_dir().join("rsx_pca_arrow_test");
        std::fs::create_dir_all(&dir).unwrap();

        let table = make_pca_test_data(&dir);
        let out_dir = dir.join("pca_file");
        std::fs::create_dir_all(&out_dir).unwrap();

        let params = PcaParams {
            markers_table_path: table.to_str().unwrap().to_string(),
            output_dir: out_dir.to_str().unwrap().to_string(),
            min_depth: 1,
            n_components: Some(3),
        };

        run(&params).unwrap();

        let eigen_file = std::fs::read_to_string(out_dir.join("eigenvalues.tsv")).unwrap();
        let eigen_lines: Vec<&str> = eigen_file
            .lines()
            .filter(|l| !l.starts_with("component"))
            .collect();
        let n_file_rows = eigen_lines.len();

        let arrow = run_to_arrow(&params).expect("pca run_to_arrow must succeed");

        assert_eq!(arrow.eigenvalues.num_rows(), n_file_rows);
        assert_eq!(arrow.loadings.num_rows(), 5);
        assert!(arrow.total_variance > 0.0);
        assert_eq!(arrow.n_components, 3);
    }
}
