//! Sankoff parsimony cost matrix for the 16-state (powerset of `{A,C,G,T}`)
//! family alphabet, and calibration of its parameters from real pairwise
//! SNP/shared counts. See `docmd/theory/evolutionary_distances.md`,
//! "Sankoff parsimony as the resolution of the 16-state model problem".

use crate::{BasePairTally, RawSnpDistanceOutput};

/// Transition base pairs in the project's fixed bit convention (see
/// `siblings::central_base`: bit 0=A, 1=C, 2=G, 3=T). All other single-bit
/// swaps are transversions.
const TRANSITION_PAIRS: [(u8, u8); 2] = [(0, 2), (1, 3)]; // A<->G, C<->T

/// Tunable weights for the 16-state Sankoff cost matrix — see "A concrete
/// Sankoff cost matrix for the 16-state alphabet" in the design doc.
///
/// Only two costs, not three: an earlier version had a separate `c_gl` for
/// "gain/loss between two nonempty states" alongside `c_ctx` for "collapse
/// to `∅`", but presence/absence tracking never observes "the flanks"
/// independently of a whole 31-mer — losing one member of a family while
/// others remain (`{A,C}->{A}`) and losing the last one (`{A}->∅`) are the
/// same event: a specific, complete, homologous 31-mer that used to be
/// observed no longer is. Both are exactly what `c_ctx` (see
/// `c_ctx_from_p_hat`) prices, so it is used for every gain/loss uniformly
/// — including a compound family losing several members at once, `∅`
/// included, which costs `|X|*c_ctx` (no flat shortcut for `∅`: see
/// `build_cost_matrix`'s docs for why an earlier version's flat special
/// case there didn't hold up).
#[derive(Debug, Clone, Copy)]
pub struct SankoffWeights {
    /// Cost of a single-base substitution `a <-> b` (`a != b`), symmetric
    /// (`sub_cost[a][b] == sub_cost[b][a]`) and only meaningful off the
    /// diagonal. The full 6-category (AC, AG, AT, CG, CT, GT) resolution a
    /// symmetric cost matrix allows — see `substitution_costs_from_tally`
    /// for calibrating it from real data, or `SankoffWeights::ts_tv` for
    /// the coarser 2-category (transition/transversion) convenience
    /// constructor.
    pub sub_cost: [[f64; 4]; 4],
    /// Cost of losing or gaining one family member — the same constant
    /// everywhere a member is gained or lost, member count and target state
    /// (`∅` or not) included: see the struct docs.
    pub c_ctx: f64,
}

impl SankoffWeights {
    /// Convenience constructor for the coarser 2-category
    /// (transition/transversion) substitution model: every transition pair
    /// (A<->G, C<->T) costs `c_ts`, every transversion pair costs `c_tv`.
    pub fn ts_tv(c_ts: f64, c_tv: f64, c_ctx: f64) -> Self {
        let mut sub_cost = [[0.0; 4]; 4];
        for a in 0..4usize {
            for b in 0..4usize {
                if a == b {
                    continue;
                }
                let is_ts = TRANSITION_PAIRS.contains(&(a as u8, b as u8));
                sub_cost[a][b] = if is_ts { c_ts } else { c_tv };
            }
        }
        Self { sub_cost, c_ctx }
    }
}

/// Build the 16x16 Sankoff step-cost matrix, states indexed `0..=15` by
/// their bitmask (matching `siblings::iupac_code`'s convention — state `0`
/// is `∅`).
///
/// All 16 states, `∅` included, form a single graph: an edge of weight
/// `c_ts`/`c_tv` between two states of equal cardinality differing by
/// exactly one element (classified by whether the swapped bases form a
/// transition or transversion pair), and an edge of weight `c_ctx` between
/// a state and one that is its strict subset plus exactly one element
/// (`∅` is a strict subset of every singleton state, so it connects
/// directly to each of them this way — no special case). `cost(X, Y)` is
/// the shortest-path distance in that graph (16 nodes — Floyd-Warshall,
/// trivial at this size), uniformly, including `X <-> ∅`: a compound
/// family losing several members at once, `∅` included, costs `|X|*c_ctx`
/// via that many single-element steps — no cheaper flat alternative for
/// `∅` specifically, since nothing distinguishes it from any other
/// multi-element loss (see design doc — an earlier flat special case here
/// was inconsistent with how every *other* multi-element transformation is
/// already priced, and had no principled justification once that was
/// noticed).
pub fn build_cost_matrix(w: &SankoffWeights) -> [[f64; 16]; 16] {
    const INF: f64 = f64::INFINITY;
    let mut dist = [[INF; 16]; 16];
    for (i, row) in dist.iter_mut().enumerate() {
        row[i] = 0.0;
    }

    for x in 0u8..16 {
        for y in (x + 1)..16 {
            let (xu, yu) = (x as usize, y as usize);
            let diff = x ^ y;
            let card_x = x.count_ones();
            let card_y = y.count_ones();
            let edge = if card_x == card_y && diff.count_ones() == 2 {
                // Exactly one element swapped: the bit only in x is the base
                // leaving, the bit only in y is the base entering.
                let a = (x & diff).trailing_zeros() as u8;
                let b = (y & diff).trailing_zeros() as u8;
                Some(w.sub_cost[a as usize][b as usize])
            } else if card_x.abs_diff(card_y) == 1 && (x & y) == x.min(y) {
                // True subset relationship: a pure gain/loss of one element.
                Some(w.c_ctx)
            } else {
                None
            };
            if let Some(cost) = edge {
                dist[xu][yu] = cost;
                dist[yu][xu] = cost;
            }
        }
    }

    for k in 0..16 {
        for i in 0..16 {
            if dist[i][k].is_infinite() {
                continue;
            }
            for j in 0..16 {
                let via = dist[i][k] + dist[k][j];
                if via < dist[i][j] {
                    dist[i][j] = via;
                }
            }
        }
    }

    dist
}

/// `p_hat` calibrated from real pairwise SNP/shared counts, restricted to
/// pairs below `ratio_ceiling`, and its variance as a pooled Bernoulli
/// proportion.
#[derive(Debug, Clone, Copy)]
pub struct PHatEstimate {
    pub p_hat: f64,
    /// `p_hat*(1-p_hat) / n_loci_total` — the pooled estimator's variance is
    /// governed by the total number of loci across all included pairs, not
    /// by any single pair's count (see design doc: this is why the
    /// exclusion criterion below is the per-pair *ratio*, not a per-pair
    /// minimum-count floor — a low-count pair barely moves the pool either
    /// way, but a saturated pair with plenty of loci would bias it).
    pub variance: f64,
    pub n_pairs_included: usize,
    pub n_loci_total: u64,
}

/// Pool `snp`/`shared` counts across all genome pairs whose per-pair ratio
/// `snp/(snp+shared)` is at most `ratio_ceiling`, then estimate
/// `p_hat = sum(snp) / sum(snp+shared)` over the included pairs.
///
/// A pair with no eligible locus at all (`snp+shared == 0`) is always
/// excluded (nothing to pool). `ratio_ceiling` should be well below 1.0 —
/// pairs at or near saturation (e.g. cross-domain comparisons, where nearly
/// every shared central-position family already differs) carry no
/// information about `p_hat` and bias it upward if pooled in.
pub fn calibrate_p_hat(raw: &RawSnpDistanceOutput, ratio_ceiling: f64) -> PHatEstimate {
    let n = raw.snp.nrows();
    let mut snp_sum: u64 = 0;
    let mut total_sum: u64 = 0;
    let mut n_pairs = 0usize;
    for i in 0..n {
        for j in (i + 1)..n {
            let snp = raw.snp[[i, j]];
            let shared = raw.shared[[i, j]];
            let total = snp + shared;
            if total == 0 {
                continue;
            }
            let ratio = snp as f64 / total as f64;
            if ratio > ratio_ceiling {
                continue;
            }
            snp_sum += snp;
            total_sum += total;
            n_pairs += 1;
        }
    }
    let p_hat = if total_sum > 0 { snp_sum as f64 / total_sum as f64 } else { 0.0 };
    let variance = if total_sum > 0 {
        p_hat * (1.0 - p_hat) / total_sum as f64
    } else {
        0.0
    };
    PHatEstimate { p_hat, variance, n_pairs_included: n_pairs, n_loci_total: total_sum }
}

/// `c_ctx(p) = mean_sub_cost * [(2m*p) / (1 - (1-p)^(2m)) + p]` — see
/// "Context, detectability, and a 3-way ordinal distance per pair" in the
/// design doc for the bracketed term's derivation. `m` is the flank length
/// on *each* side of the central position (`k = 2m+1`); pass `(k-1)/2`, not
/// the project's minimizer-size parameter of the same name.
///
/// The bracketed term is `E[mutation count | at least one occurred]` — a
/// *count* of mutations, not a cost. An earlier version used it directly as
/// the cost, implicitly pricing every mutation at a flat `1` regardless of
/// type. That stopped making sense once substitution costs were calibrated
/// per base-pair type (`substitution_costs_from_tally`): transitions and
/// transversions aren't equally likely (roughly 5-6x apart in this
/// project's own data) nor equally costly, so "one mutation" isn't worth a
/// flat unit — it's worth whatever the *empirical mix* of mutation types is
/// worth on average. `mean_sub_cost` (see `mean_substitution_cost`) is that
/// weighted average, turning the expected mutation *count* into an actual
/// expected cost.
pub fn c_ctx_from_p_hat(p_hat: f64, m: usize, mean_sub_cost: f64) -> f64 {
    if p_hat <= 0.0 {
        return 0.0;
    }
    let two_m = 2.0 * m as f64;
    let expected_mutations = (two_m * p_hat) / (1.0 - (1.0 - p_hat).powf(two_m)) + p_hat;
    expected_mutations * mean_sub_cost
}

/// Mean substitution cost, weighted by each of the 6 base-pair categories'
/// observed frequency in `tally` — the empirical average cost of "one
/// mutation" under the calibrated substitution spectrum. Falls back to
/// `1.0` (a flat per-mutation cost) if `tally` has no signal at all (no
/// substitution ever observed), rather than dividing by zero.
pub fn mean_substitution_cost(tally: &BasePairTally, sub_cost: &[[f64; 4]; 4]) -> f64 {
    let mut weighted_sum = 0.0;
    let mut total = 0u64;
    for a in 0..4 {
        for b in (a + 1)..4 {
            let count = tally.counts[a][b];
            weighted_sum += count as f64 * sub_cost[a][b];
            total += count;
        }
    }
    if total == 0 {
        return 1.0;
    }
    weighted_sum / total as f64
}

/// Derive the 4x4 symmetric substitution cost table from an observed
/// [`BasePairTally`]: `cost(a,b) = -ln(rate(a,b))`, normalised so the most
/// frequently observed substitution type costs exactly `1.0` — the standard
/// generalised-parsimony step-weighting heuristic (see design doc,
/// "Transition/transversion refinement"), generalised here from 2
/// categories (Ts/Tv) to the full 6-category resolution a symmetric matrix
/// allows: a rarer substitution type is treated as less parsimonious to
/// invoke, and therefore costs more.
///
/// A pair type never observed at all (`counts[a][b] == 0`) gets an infinite
/// cost — parsimony should never spend a mutation on something with no
/// empirical support in this data.
pub fn substitution_costs_from_tally(tally: &BasePairTally) -> [[f64; 4]; 4] {
    let total: u64 = (0..4)
        .flat_map(|a| (a + 1..4).map(move |b| (a, b)))
        .map(|(a, b)| tally.counts[a][b])
        .sum();

    let mut raw = [[0.0f64; 4]; 4];
    let mut min_cost = f64::INFINITY;
    for a in 0..4 {
        for b in (a + 1)..4 {
            let count = tally.counts[a][b];
            let cost = if count > 0 && total > 0 {
                -((count as f64 / total as f64).ln())
            } else {
                f64::INFINITY
            };
            raw[a][b] = cost;
            raw[b][a] = cost;
            if cost < min_cost {
                min_cost = cost;
            }
        }
    }

    let mut out = [[0.0f64; 4]; 4];
    for a in 0..4 {
        for b in 0..4 {
            if a != b {
                out[a][b] = raw[a][b] / min_cost;
            }
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::Array2;

    fn weights() -> SankoffWeights {
        SankoffWeights::ts_tv(1.0, 2.0, 10.0)
    }

    #[test]
    fn identity_is_zero() {
        let m = build_cost_matrix(&weights());
        for row in m.iter().enumerate() {
            assert_eq!(m[row.0][row.0], 0.0);
        }
    }

    #[test]
    fn single_base_transition_costs_c_ts() {
        let m = build_cost_matrix(&weights());
        // A = 0b0001 (state 1), G = 0b0100 (state 4): A<->G is a transition.
        assert_eq!(m[1][4], 1.0);
        assert_eq!(m[4][1], 1.0);
    }

    #[test]
    fn single_base_transversion_costs_c_tv() {
        let m = build_cost_matrix(&weights());
        // A = 0b0001 (state 1), C = 0b0010 (state 2): A<->C is a transversion.
        assert_eq!(m[1][2], 2.0);
    }

    #[test]
    fn gain_loss_costs_c_ctx() {
        let m = build_cost_matrix(&weights());
        // A = 0b0001 (state 1) -> {A,C} = 0b0011 (state 3): pure gain of C.
        assert_eq!(m[1][3], 10.0);
    }

    #[test]
    fn empty_set_costs_scale_with_cardinality() {
        let m = build_cost_matrix(&weights());
        assert_eq!(m[0][1], 10.0); // {A} -> ∅: one member, one c_ctx step.
        // N (all four) -> ∅: no flat shortcut — four single-element steps,
        // same as losing four members down to a nonempty state would cost.
        assert_eq!(m[0][15], 40.0);
        assert_eq!(m[0][0], 0.0);
    }

    #[test]
    fn matrix_is_symmetric() {
        let m = build_cost_matrix(&weights());
        for i in 0..16 {
            for j in 0..16 {
                assert_eq!(m[i][j], m[j][i], "asymmetry at ({i},{j})");
            }
        }
    }

    #[test]
    fn calibrate_p_hat_excludes_saturated_pairs() {
        // 3 genomes: pair (0,1) informative (10% SNP), pair (0,2) saturated
        // (100% SNP) — must be excluded from the pooled estimate.
        let mut snp = Array2::<u64>::zeros((3, 3));
        let mut shared = Array2::<u64>::zeros((3, 3));
        snp[[0, 1]] = 10;
        shared[[0, 1]] = 90;
        snp[[1, 0]] = 10;
        shared[[1, 0]] = 90;
        snp[[0, 2]] = 50;
        shared[[0, 2]] = 0;
        snp[[2, 0]] = 50;
        shared[[2, 0]] = 0;
        let raw = RawSnpDistanceOutput { snp, shared };

        let est = calibrate_p_hat(&raw, 0.5);
        assert_eq!(est.n_pairs_included, 1);
        assert_eq!(est.n_loci_total, 100);
        assert!((est.p_hat - 0.1).abs() < 1e-9);
    }

    #[test]
    fn c_ctx_from_p_hat_zero_is_zero() {
        assert_eq!(c_ctx_from_p_hat(0.0, 15, 1.5), 0.0);
    }

    #[test]
    fn c_ctx_from_p_hat_matches_known_value() {
        // p=0.1, m=15: [(30*0.1)/(1-(0.9)^30) + 0.1] * mean_sub_cost
        let m = 15usize;
        let p = 0.1;
        let mean_sub_cost = 1.5;
        let expected = ((2.0 * m as f64 * p) / (1.0 - (1.0 - p).powf(2.0 * m as f64)) + p) * mean_sub_cost;
        assert!((c_ctx_from_p_hat(p, m, mean_sub_cost) - expected).abs() < 1e-12);
    }

    #[test]
    fn mean_substitution_cost_is_weighted_average() {
        let mut counts = [[0u64; 4]; 4];
        counts[0][1] = 3; // A/C
        counts[1][0] = 3;
        counts[0][2] = 1; // A/G
        counts[2][0] = 1;
        let tally = BasePairTally { counts };
        let mut sub_cost = [[0.0f64; 4]; 4];
        sub_cost[0][1] = 2.0;
        sub_cost[1][0] = 2.0;
        sub_cost[0][2] = 1.0;
        sub_cost[2][0] = 1.0;
        // (3*2.0 + 1*1.0) / 4 = 1.75
        assert!((mean_substitution_cost(&tally, &sub_cost) - 1.75).abs() < 1e-12);
    }

    #[test]
    fn mean_substitution_cost_falls_back_to_one_with_no_data() {
        let tally = BasePairTally { counts: [[0u64; 4]; 4] };
        let sub_cost = [[0.0f64; 4]; 4];
        assert_eq!(mean_substitution_cost(&tally, &sub_cost), 1.0);
    }
}
