use std::{fs, io, path::{Path, PathBuf}};

use ndarray::{Array1, Array2};
use rayon::prelude::*;

use crate::builder::PersistentCompactIntVecBuilder;
use crate::meta::MatrixMeta;
use crate::reader::PersistentCompactIntVec;

fn col_path(dir: &Path, col: usize) -> PathBuf {
    dir.join(format!("col_{col:06}.pciv"))
}

pub struct PersistentCompactIntMatrix {
    cols: Vec<PersistentCompactIntVec>,
    n:    usize,
}

impl PersistentCompactIntMatrix {
    pub fn open(dir: &Path) -> io::Result<Self> {
        let meta = MatrixMeta::load(dir)?;
        let cols = (0..meta.n_cols)
            .map(|c| PersistentCompactIntVec::open(&col_path(dir, c)))
            .collect::<io::Result<Vec<_>>>()?;
        Ok(Self { cols, n: meta.n })
    }

    pub fn n(&self)      -> usize { self.n }
    pub fn n_cols(&self) -> usize { self.cols.len() }
    pub fn col(&self, c: usize) -> &PersistentCompactIntVec { &self.cols[c] }

    pub fn row(&self, slot: usize) -> Box<[u32]> {
        self.cols.iter().map(|c| c.get(slot)).collect()
    }

    // ── Distance matrices ─────────────────────────────────────────────────────

    pub fn bray_dist_matrix(&self) -> Array2<f64> {
        let sum_min  = self.partial_bray_dist_matrix();
        let col_sums = self.sum();
        let n = self.n_cols();
        let mut m = Array2::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                if i != j {
                    let denom = col_sums[i] + col_sums[j];
                    m[[i, j]] = if denom == 0 { 0.0 }
                                 else { 1.0 - 2.0 * sum_min[[i, j]] as f64 / denom as f64 };
                }
            }
        }
        m
    }

    pub fn relfreq_bray_dist_matrix(&self) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).relfreq_bray_dist(self.col(j)))
    }

    pub fn euclidean_dist_matrix(&self) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).euclidean_dist(self.col(j)))
    }

    pub fn relfreq_euclidean_dist_matrix(&self) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).relfreq_euclidean_dist(self.col(j)))
    }

    pub fn hellinger_dist_matrix(&self) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).hellinger_dist(self.col(j)))
    }

    pub fn jaccard_dist_matrix(&self) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).jaccard_dist(self.col(j)))
    }

    pub fn threshold_jaccard_dist_matrix(&self, threshold: u32) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).threshold_jaccard_dist(self.col(j), threshold))
    }

    /// Returns the sum of each column as `Array1<u64>`.
    pub fn sum(&self) -> Array1<u64> {
        let sums: Vec<u64> = (0..self.n_cols())
            .into_par_iter()
            .map(|c| self.col(c).sum())
            .collect();
        Array1::from_vec(sums)
    }

    // ── Partial matrices (additively decomposable across layers) ──────────────

    /// Returns `sum_min[n×n]` where `sum_min[i,j]` = Σ_slot min(col_i[slot], col_j[slot]).
    /// The denominator `col_sums[i] + col_sums[j]` is obtained from `self.sum()`.
    /// Additive across layers by element-wise addition.
    pub fn partial_bray_dist_matrix(&self) -> Array2<u64> {
        self.pairwise_u64(|i, j| self.col(i).partial_bray_dist(self.col(j)))
    }

    /// Returns sum of squared differences `[n×n]`.
    /// Reduce across layers by element-wise addition, then take `sqrt` for the final distance.
    pub fn partial_euclidean_dist_matrix(&self) -> Array2<f64> {
        self.pairwise(|i, j| self.col(i).partial_euclidean_dist(self.col(j)))
    }

    /// Returns `(inter[n×n], union[n×n])` for threshold-Jaccard.
    /// Reduce across layers by element-wise addition before computing `1 - inter/union`.
    pub fn partial_threshold_jaccard_dist_matrix(
        &self,
        threshold: u32,
    ) -> (Array2<u64>, Array2<u64>) {
        let n = self.n_cols();
        let pairs = upper_pairs(n);

        let results: Vec<(usize, usize, u64, u64)> = pairs
            .into_par_iter()
            .map(|(i, j)| {
                let (inter, union) =
                    self.col(i).partial_threshold_jaccard_dist(self.col(j), threshold);
                (i, j, inter, union)
            })
            .collect();

        let mut inter_m = Array2::zeros((n, n));
        let mut union_m = Array2::zeros((n, n));
        for (i, j, inter, union) in results {
            inter_m[[i, j]] = inter; inter_m[[j, i]] = inter;
            union_m[[i, j]] = union; union_m[[j, i]] = union;
        }
        (inter_m, union_m)
    }

    /// Returns matrix of `Σ_slot min(col_i[slot]/sum_i, col_j[slot]/sum_j)` per pair.
    /// `col_sums` must be the GLOBAL sums across all layers/partitions.
    /// Reduce across layers by element-wise addition; final distance = `1 - value`.
    pub fn partial_relfreq_bray_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
        self.pairwise(|i, j| {
            self.col(i).partial_relfreq_bray_dist(
                self.col(j),
                col_sums[i] as f64,
                col_sums[j] as f64,
            )
        })
    }

    /// Returns matrix of `Σ_slot (col_i[slot]/sum_i - col_j[slot]/sum_j)²` per pair.
    /// `col_sums` must be the GLOBAL sums across all layers/partitions.
    /// Reduce across layers by element-wise addition; final distance = `sqrt(value)`.
    pub fn partial_relfreq_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
        self.pairwise(|i, j| {
            self.col(i).partial_relfreq_euclidean_dist(
                self.col(j),
                col_sums[i] as f64,
                col_sums[j] as f64,
            )
        })
    }

    /// Returns matrix of `Σ_slot (√(col_i/sum_i) - √(col_j/sum_j))²` per pair.
    /// `col_sums` must be the GLOBAL sums across all layers/partitions.
    /// Reduce across layers by element-wise addition; final distance = `sqrt(value) / √2`.
    pub fn partial_hellinger_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
        self.pairwise(|i, j| {
            self.col(i).partial_hellinger_euclidean_dist(
                self.col(j),
                col_sums[i] as f64,
                col_sums[j] as f64,
            )
        })
    }

    // ── Private helpers ───────────────────────────────────────────────────────

    fn pairwise(&self, f: impl Fn(usize, usize) -> f64 + Sync) -> Array2<f64> {
        let n = self.n_cols();
        let results: Vec<(usize, usize, f64)> = upper_pairs(n)
            .into_par_iter()
            .map(|(i, j)| (i, j, f(i, j)))
            .collect();
        fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
    }

    fn pairwise_u64(&self, f: impl Fn(usize, usize) -> u64 + Sync) -> Array2<u64> {
        let n = self.n_cols();
        let results: Vec<(usize, usize, u64)> = upper_pairs(n)
            .into_par_iter()
            .map(|(i, j)| (i, j, f(i, j)))
            .collect();
        fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
    }
}

fn upper_pairs(n: usize) -> Vec<(usize, usize)> {
    (0..n).flat_map(|i| (i + 1..n).map(move |j| (i, j))).collect()
}

fn fill_symmetric<T>(n: usize, vals: impl Iterator<Item = (usize, usize, T, T)>) -> Array2<T>
where
    T: Clone + Default,
{
    let mut m = Array2::from_elem((n, n), T::default());
    for (i, j, vij, vji) in vals {
        m[[i, j]] = vij;
        m[[j, i]] = vji;
    }
    m
}

// ── Column append ─────────────────────────────────────────────────────────────

impl PersistentCompactIntMatrix {
    /// Append a new column to an existing matrix on disk.
    ///
    /// Reads `meta.json` to obtain `n` and the current column count, writes
    /// `col_{n_cols:06}.pciv` filled by `value_of(slot)`, then increments
    /// `n_cols` in `meta.json`.
    pub fn append_column(dir: &Path, value_of: impl Fn(usize) -> u32) -> io::Result<()> {
        let mut meta = MatrixMeta::load(dir)?;
        let mut b = PersistentCompactIntVecBuilder::new(meta.n, &col_path(dir, meta.n_cols))?;
        for slot in 0..meta.n {
            b.set(slot, value_of(slot));
        }
        b.close()?;
        meta.n_cols += 1;
        meta.save(dir)
    }
}

// ── Trait impls ───────────────────────────────────────────────────────────────

use crate::traits::{ColumnWeights, CountPartials};

impl ColumnWeights for PersistentCompactIntMatrix {
    fn col_weights(&self) -> Array1<u64> { self.sum() }
}

impl CountPartials for PersistentCompactIntMatrix {
    fn partial_bray(&self) -> Array2<u64> {
        self.partial_bray_dist_matrix()
    }
    fn partial_euclidean(&self) -> Array2<f64> {
        self.partial_euclidean_dist_matrix()
    }
    fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
        self.partial_threshold_jaccard_dist_matrix(threshold)
    }
    fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
        self.partial_relfreq_bray_dist_matrix(global)
    }
    fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
        self.partial_relfreq_euclidean_dist_matrix(global)
    }
    fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
        self.partial_hellinger_euclidean_dist_matrix(global)
    }
}

// ── Builder ───────────────────────────────────────────────────────────────────

pub struct PersistentCompactIntMatrixBuilder {
    dir:    PathBuf,
    n:      usize,
    n_cols: usize,
}

impl PersistentCompactIntMatrixBuilder {
    pub fn new(n: usize, dir: &Path) -> io::Result<Self> {
        fs::create_dir_all(dir)?;
        Ok(Self { dir: dir.to_path_buf(), n, n_cols: 0 })
    }

    pub fn n(&self)      -> usize { self.n }
    pub fn n_cols(&self) -> usize { self.n_cols }

    pub fn add_col(&mut self) -> io::Result<PersistentCompactIntVecBuilder> {
        let path = col_path(&self.dir, self.n_cols);
        self.n_cols += 1;
        PersistentCompactIntVecBuilder::new(self.n, &path)
    }

    pub fn close(self) -> io::Result<()> {
        MatrixMeta { n: self.n, n_cols: self.n_cols }.save(&self.dir)
    }
}
