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)]
#[path = "tests/layered_store.rs"]
mod tests;
