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

use obicompactvec::traits::{BitPartials, ColumnWeights, CountPartials};

/// A store that aggregates a `Vec<S>` — one entry per layer (within a partition)
/// or one entry per partition.
///
/// Blanket impls of `ColumnWeights`, `CountPartials`, and `BitPartials` propagate
/// automatically: `LayeredStore<LayeredStore<S>>` implements the same traits as
/// `LayeredStore<S>`, giving the partitioned level for free.
pub struct LayeredStore<S>(pub Vec<S>);

impl<S> LayeredStore<S> {
    pub fn new(layers: Vec<S>) -> Self { Self(layers) }
    pub fn layers(&self) -> &[S] { &self.0 }
    pub fn n_layers(&self) -> usize { self.0.len() }
    pub fn is_empty(&self) -> bool { self.0.is_empty() }
}

// ── ColumnWeights ─────────────────────────────────────────────────────────────

impl<S: ColumnWeights> ColumnWeights for LayeredStore<S> {
    fn col_weights(&self) -> Array1<u64> {
        self.0.par_iter()
            .map(|s| s.col_weights())
            .reduce_with(|a, b| a + b)
            .unwrap_or_else(|| Array1::zeros(0))
    }
}

// ── CountPartials ─────────────────────────────────────────────────────────────

impl<S: CountPartials> CountPartials for LayeredStore<S> {
    fn partial_bray(&self) -> Array2<u64> {
        self.0.par_iter()
            .map(|s| s.partial_bray())
            .reduce_with(|a, b| a + b)
            .unwrap()
    }

    fn partial_euclidean(&self) -> Array2<f64> {
        self.0.par_iter()
            .map(|s| s.partial_euclidean())
            .reduce_with(|a, b| a + b)
            .unwrap()
    }

    fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
        self.0.par_iter()
            .map(|s| s.partial_threshold_jaccard(threshold))
            .reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu))
            .unwrap()
    }

    fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
        self.0.par_iter()
            .map(|s| s.partial_relfreq_bray(global))
            .reduce_with(|a, b| a + b)
            .unwrap()
    }

    fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
        self.0.par_iter()
            .map(|s| s.partial_relfreq_euclidean(global))
            .reduce_with(|a, b| a + b)
            .unwrap()
    }

    fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
        self.0.par_iter()
            .map(|s| s.partial_hellinger(global))
            .reduce_with(|a, b| a + b)
            .unwrap()
    }
}

// ── BitPartials ───────────────────────────────────────────────────────────────

impl<S: BitPartials> BitPartials for LayeredStore<S> {
    fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
        self.0.par_iter()
            .map(|s| s.partial_jaccard())
            .reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu))
            .unwrap()
    }

    fn partial_hamming(&self) -> Array2<u64> {
        self.0.par_iter()
            .map(|s| s.partial_hamming())
            .reduce_with(|a, b| a + b)
            .unwrap()
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use obicompactvec::{
        PersistentBitMatrix, PersistentBitMatrixBuilder,
        PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
    };
    use tempfile::tempdir;

    fn make_int_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
        let n = cols.first().map_or(0, |c| c.len());
        let dir = tempdir().unwrap();
        let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
        for &col in cols {
            let mut cb = b.add_col().unwrap();
            for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
            cb.close().unwrap();
        }
        b.close().unwrap();
        let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
        (dir, m)
    }

    fn make_bit_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
        let n = cols.first().map_or(0, |c| c.len());
        let dir = tempdir().unwrap();
        let mut b = PersistentBitMatrixBuilder::new(n, &dir.path().join("presence")).unwrap();
        for &col in cols {
            let mut cb = b.add_col().unwrap();
            for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
            cb.close().unwrap();
        }
        b.close().unwrap();
        let m = PersistentBitMatrix::open(dir.path()).unwrap();
        (dir, m)
    }

    // ── ColumnWeights ─────────────────────────────────────────────────────────

    #[test]
    fn col_weights_sums_across_layers() {
        // layer 0: col0=[1,2], col1=[3,4]   → weights [3, 7]
        // layer 1: col0=[10,0], col1=[0,10]  → weights [10, 10]
        // combined: [13, 17]
        let (_d0, m0) = make_int_matrix(&[&[1, 2], &[3, 4]]);
        let (_d1, m1) = make_int_matrix(&[&[10, 0], &[0, 10]]);
        let store = LayeredStore::new(vec![m0, m1]);
        let w = store.col_weights();
        assert_eq!(w[0], 13);
        assert_eq!(w[1], 17);
    }

    #[test]
    fn col_weights_bit_sums_across_layers() {
        // layer 0: col0=[T,F,T], col1=[F,T,T]  → counts [2, 2]
        // layer 1: col0=[F,F,T], col1=[T,T,F]  → counts [1, 2]
        // combined: [3, 4]
        let (_d0, m0) = make_bit_matrix(&[&[true, false, true], &[false, true, true]]);
        let (_d1, m1) = make_bit_matrix(&[&[false, false, true], &[true, true, false]]);
        let store = LayeredStore::new(vec![m0, m1]);
        let w = store.col_weights();
        assert_eq!(w[0], 3);
        assert_eq!(w[1], 4);
    }

    // ── CountPartials — layered (one partition) ───────────────────────────────

    #[test]
    fn layered_bray_matches_combined() {
        // Split [1,2,3,4,5] across two layers; bray dist should equal direct computation
        // on [1,2,3,4,5] for each column pair.
        // col0=[1,2,3,4,5], col1=[5,4,3,2,1]
        let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]);     // slots 0-1
        let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]); // slots 2-4
        let store = LayeredStore::new(vec![m0, m1]);

        // direct on full data
        let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
        let expected = CountPartials::bray_dist_matrix(&mf);
        let got      = CountPartials::bray_dist_matrix(&store);
        assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "bray [0,1]");
        assert!((got[[1, 0]] - expected[[1, 0]]).abs() < 1e-12, "bray [1,0]");
    }

    #[test]
    fn layered_relfreq_bray_matches_combined() {
        let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]);
        let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
        let store = LayeredStore::new(vec![m0, m1]);

        let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
        let expected = CountPartials::relfreq_bray_dist_matrix(&mf);
        let got      = CountPartials::relfreq_bray_dist_matrix(&store);
        assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "relfreq_bray [0,1]");
    }

    #[test]
    fn layered_euclidean_matches_combined() {
        let (_d0, m0) = make_int_matrix(&[&[3, 0], &[0, 4]]);
        let (_d1, m1) = make_int_matrix(&[&[1, 1], &[2, 2]]);
        let store = LayeredStore::new(vec![m0, m1]);

        let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 4, 2, 2]]);
        let expected = CountPartials::euclidean_dist_matrix(&mf);
        let got      = CountPartials::euclidean_dist_matrix(&store);
        assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "euclidean [0,1]");
    }

    // ── CountPartials — partitioned (LayeredStore<LayeredStore<_>>) ───────────

    #[test]
    fn partitioned_bray_matches_combined() {
        // partition 0: slots [1,2,3,4,5]  col0 vs col1
        // partition 1: slots [10,20]      col0 vs col1
        let (_d0, p0) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
        let (_d1, p1) = make_int_matrix(&[&[10, 20], &[20, 10]]);

        let partitioned = LayeredStore::new(vec![
            LayeredStore::new(vec![p0]),
            LayeredStore::new(vec![p1]),
        ]);

        let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5, 10, 20], &[5, 4, 3, 2, 1, 20, 10]]);
        let expected = CountPartials::bray_dist_matrix(&mf);
        let got      = CountPartials::bray_dist_matrix(&partitioned);
        assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "partitioned bray [0,1]");
    }

    // ── BitPartials ───────────────────────────────────────────────────────────

    #[test]
    fn layered_jaccard_matches_combined() {
        let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
        let (_d1, m1) = make_bit_matrix(&[&[true, true], &[true, false]]);
        let store = LayeredStore::new(vec![m0, m1]);

        let (_df, mf) = make_bit_matrix(&[
            &[true, false, true, true],
            &[false, true, true, false],
        ]);
        let expected = BitPartials::jaccard_dist_matrix(&mf);
        let got      = BitPartials::jaccard_dist_matrix(&store);
        assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "jaccard [0,1]");
    }

    #[test]
    fn layered_hamming_matches_combined() {
        let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
        let (_d1, m1) = make_bit_matrix(&[&[true, true], &[false, false]]);
        let store = LayeredStore::new(vec![m0, m1]);

        let (_df, mf) = make_bit_matrix(&[
            &[true, false, true, true],
            &[false, true, false, false],
        ]);
        let expected = BitPartials::hamming_dist_matrix(&mf);
        let got      = BitPartials::hamming_dist_matrix(&store);
        assert_eq!(got[[0, 1]], expected[[0, 1]], "hamming [0,1]");
    }
}
