use ndarray::{Array1, Array2};

/// Convert a Jaccard distance matrix (`1 - J`) into a Mash distance matrix, per
/// https://mash.readthedocs.io/en/latest/distances.html:
/// `D = -1/k * ln(2J / (1+J))`.
fn jaccard_to_mash(d_jaccard: &Array2<f64>, k: usize) -> Array2<f64> {
    d_jaccard.mapv(|d| {
        let j = 1.0 - d;
        if j <= 0.0 { 1.0 }
        else { -1.0 / k as f64 * (2.0 * j / (1.0 + j)).ln() }
    })
}

// ── Column-level weight statistic — total count or presence count per column.
/// Additive across layers and partitions; used as denominator in normalised distances.
///
/// `partial_kmer_counts` returns the number of **distinct k-mers** present per
/// column (presence = 1 entries; count > 0 entries). For presence matrices this
/// equals `col_weights`; for count matrices it differs (count_nonzero vs sum).
pub trait ColumnWeights: Send + Sync {
    fn col_weights(&self) -> Array1<u64>;

    fn partial_kmer_counts(&self) -> Array1<u64> {
        self.col_weights()
    }
}

/// Partial distance matrices for count-based data (`PersistentCompactIntMatrix`).
///
/// Every `partial_*` method returns an additive component: element-wise summing the results
/// across layers then across partitions yields the global partial, from which the final
/// distance is computed via the corresponding provided method.
///
/// Normalised methods (`partial_relfreq_*`, `partial_hellinger`) require the **global**
/// `col_weights` (summed across all layers and partitions) as a parameter.  The provided
/// finalisation methods compute this in a first pass via `self.col_weights()`.
pub trait CountPartials: ColumnWeights {
    fn partial_bray(&self) -> Array2<u64>;
    fn partial_euclidean(&self) -> Array2<f64>;
    fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>);
    fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64>;
    fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64>;
    fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64>;

    // ── Provided finalisation methods ─────────────────────────────────────────

    fn bray_dist_matrix(&self) -> Array2<f64> {
        let sum_min = self.partial_bray();
        let w = self.col_weights();
        let n = w.len();
        let mut m = Array2::<f64>::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                if i != j {
                    let d = w[i] + w[j];
                    m[[i, j]] = if d == 0 { 0.0 }
                                 else { 1.0 - 2.0 * sum_min[[i, j]] as f64 / d as f64 };
                }
            }
        }
        m
    }

    fn euclidean_dist_matrix(&self) -> Array2<f64> {
        self.partial_euclidean().mapv(|v| v.sqrt())
    }

    fn jaccard_dist_matrix(&self) -> Array2<f64> {
        self.threshold_jaccard_dist_matrix(1)
    }

    fn threshold_jaccard_dist_matrix(&self, threshold: u32) -> Array2<f64> {
        let (inter, union) = self.partial_threshold_jaccard(threshold);
        let n = inter.shape()[0];
        let mut m = Array2::<f64>::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                if i != j {
                    let u = union[[i, j]];
                    m[[i, j]] = if u == 0 { 0.0 }
                                 else { 1.0 - inter[[i, j]] as f64 / u as f64 };
                }
            }
        }
        m
    }

    /// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived
    /// from the presence-threshold Jaccard distance.
    fn threshold_mash_dist_matrix(&self, k: usize, threshold: u32) -> Array2<f64> {
        jaccard_to_mash(&self.threshold_jaccard_dist_matrix(threshold), k)
    }

    fn relfreq_bray_dist_matrix(&self) -> Array2<f64> {
        let global = self.col_weights();
        let mut m = self.partial_relfreq_bray(&global).mapv(|v| 1.0 - v);
        let n = m.shape()[0];
        for i in 0..n { m[[i, i]] = 0.0; }
        m
    }

    fn relfreq_euclidean_dist_matrix(&self) -> Array2<f64> {
        let global = self.col_weights();
        self.partial_relfreq_euclidean(&global).mapv(|v| v.sqrt())
    }

    fn hellinger_dist_matrix(&self) -> Array2<f64> {
        let global = self.col_weights();
        let sq2 = std::f64::consts::SQRT_2;
        self.partial_hellinger(&global).mapv(|v| v.sqrt() / sq2)
    }

    /// Euclidean distance in the Hellinger (√relative-frequency) space.
    /// Equal to √2 × hellinger_dist — unnormalised variant.
    fn hellinger_euclidean_dist_matrix(&self) -> Array2<f64> {
        let global = self.col_weights();
        self.partial_hellinger(&global).mapv(|v| v.sqrt())
    }
}

/// Partial distance matrices for bit-based data (`PersistentBitMatrix`).
///
/// Both `partial_*` methods are additively decomposable across layers and partitions.
pub trait BitPartials: ColumnWeights {
    fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>);
    fn partial_hamming(&self) -> Array2<u64>;

    // ── Provided finalisation methods ─────────────────────────────────────────

    fn jaccard_dist_matrix(&self) -> Array2<f64> {
        let (inter, union) = self.partial_jaccard();
        let n = inter.shape()[0];
        let mut m = Array2::<f64>::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                if i != j {
                    let u = union[[i, j]];
                    m[[i, j]] = if u == 0 { 0.0 }
                                 else { 1.0 - inter[[i, j]] as f64 / u as f64 };
                }
            }
        }
        m
    }

    /// Mash distance (https://mash.readthedocs.io/en/latest/distances.html), derived
    /// from the Jaccard distance.
    fn mash_dist_matrix(&self, k: usize) -> Array2<f64> {
        jaccard_to_mash(&self.jaccard_dist_matrix(), k)
    }

    fn hamming_dist_matrix(&self) -> Array2<u64> {
        self.partial_hamming()
    }
}
