mtnoc-rs 0.1.0

//! Pure-Rust port of the mtNoC phylogenetic mixture-EM contributor-count core (src/mtnoc/em.py).
//!
//! Numerics match em.py's MLE path exactly (deterministic raw estimate). The bootstrap-stability
//! range uses a seeded RNG that cannot bit-match numpy's Mersenne Twister, so the bootstrap lower
//! bound is validated only within +/-1 of the Python cohort (as specified).
//!
//! Reference model (see em.py): expected allele per (haplogroup, informative position) = rCRS base
//! unless the haplogroup's cumulative phylo profile overrides it. Per read (its alleles at
//! informative positions) covering n informative sites with `mm` mismatches vs haplogroup h:
//!   logL[h] = n*ln(1-eps) + mm*(ln(eps/3)-ln(1-eps))   (eps=0.005)

use anyhow::{Context, Result};
use rustc_hash::{FxHashMap, FxHashSet};
use std::io::BufRead;

/// The compact phylogenetic reference loaded from the Python dump (dump_ref_for_rust.py).
pub struct RefData {
    pub haplo: Vec<String>,               // index-ordered haplogroup labels
    pub macro_of: Vec<String>,            // macro-clade key per haplogroup (from em.py macro())
    #[allow(dead_code)]
    pub major_of: Vec<String>,            // major-haplogroup key per haplogroup (parsed; MAJOR section)
    #[allow(dead_code)]
    pub info: Vec<i64>,                   // sorted informative positions (1-based)
    pub info_idx: FxHashMap<i64, usize>,  // position -> informative index
    /// Dense expected-allele codes, layout exp[j*H + h] (i8): -1 (unknown) or 0..3 (A/C/G/T).
    pub exp: Vec<i8>,
    pub h: usize, // number of haplogroups
    #[allow(dead_code)]
    pub m: usize, // number of informative positions
    pub lineage: Vec<Vec<u32>>,           // per haplogroup: root-first lineage as indices
    pub lineage_set: Vec<FxHashSet<u32>>, // per haplogroup: lineage members (for is_ancestor)
}

impl RefData {
    /// Parse a ref dump from any reader (e.g. an in-memory decompressed embedded asset).
    pub fn load_reader<R: BufRead>(reader: R) -> Result<Self> {
        let mut lines = reader.lines();
        let mut next = || -> Result<String> {
            lines
                .next()
                .context("unexpected EOF in ref dump")?
                .context("read line")
        };
        // HAPLO <n>
        let hdr = next()?;
        let h: usize = hdr
            .strip_prefix("HAPLO ")
            .context("expected HAPLO header")?
            .trim()
            .parse()?;
        let mut haplo = Vec::with_capacity(h);
        for _ in 0..h {
            haplo.push(next()?);
        }
        // MACRO
        anyhow::ensure!(next()?.trim() == "MACRO", "expected MACRO section");
        let mut macro_of = Vec::with_capacity(h);
        for _ in 0..h {
            macro_of.push(next()?);
        }
        // MAJOR
        anyhow::ensure!(next()?.trim() == "MAJOR", "expected MAJOR section");
        let mut major_of = Vec::with_capacity(h);
        for _ in 0..h {
            major_of.push(next()?);
        }
        // INFO <m>
        let ih = next()?;
        let m: usize = ih
            .strip_prefix("INFO ")
            .context("expected INFO header")?
            .trim()
            .parse()?;
        let mut info = Vec::with_capacity(m);
        let mut info_idx = FxHashMap::default();
        let mut rcrs_code = Vec::with_capacity(m);
        for j in 0..m {
            let l = next()?;
            let mut it = l.split_whitespace();
            let pos: i64 = it.next().context("info pos")?.parse()?;
            let c: i8 = it.next().context("info code")?.parse()?;
            info.push(pos);
            info_idx.insert(pos, j);
            rcrs_code.push(c);
        }
        // dense exp: default rcrs code per info row
        let mut exp = vec![0i8; m * h];
        for j in 0..m {
            let base = j * h;
            let c = rcrs_code[j];
            for x in &mut exp[base..base + h] {
                *x = c;
            }
        }
        // PROF <n>
        let ph = next()?;
        let np: usize = ph
            .strip_prefix("PROF ")
            .context("expected PROF header")?
            .trim()
            .parse()?;
        for _ in 0..np {
            let l = next()?;
            let mut it = l.split_whitespace();
            let hi: usize = it.next().context("prof hidx")?.parse()?;
            let k: usize = it.next().context("prof k")?.parse()?;
            for _ in 0..k {
                let j: usize = it.next().context("prof j")?.parse()?;
                let c: i8 = it.next().context("prof c")?.parse()?;
                exp[j * h + hi] = c;
            }
        }
        // LIN <n>
        let lh = next()?;
        let nl: usize = lh
            .strip_prefix("LIN ")
            .context("expected LIN header")?
            .trim()
            .parse()?;
        let mut lineage = vec![Vec::new(); h];
        let mut lineage_set = vec![FxHashSet::default(); h];
        for _ in 0..nl {
            let l = next()?;
            let mut it = l.split_whitespace();
            let hi: usize = it.next().context("lin hidx")?.parse()?;
            let k: usize = it.next().context("lin k")?.parse()?;
            let mut v = Vec::with_capacity(k);
            let mut s = FxHashSet::default();
            for _ in 0..k {
                let x: u32 = it.next().context("lin idx")?.parse()?;
                v.push(x);
                s.insert(x);
            }
            lineage[hi] = v;
            lineage_set[hi] = s;
        }
        Ok(RefData {
            haplo,
            macro_of,
            major_of,
            info,
            info_idx,
            exp,
            h,
            m,
            lineage,
            lineage_set,
        })
    }

    /// is_ancestor(a, d): true if haplogroup index `a` is on the lineage of `d` (or equals it).
    #[inline]
    fn is_ancestor(&self, a: u32, d: u32) -> bool {
        self.lineage_set[d as usize].contains(&a)
    }
}

/// A per-read observation: list of (informative-position index, allele code 0..3).
pub type ReadProfile = Vec<(usize, i8)>;

/// RNG self-check: prints values to compare against numpy RandomState (dev/validation only).
pub fn rng_selftest() {
    let mut r = Mt19937::new(0);
    let u: Vec<u32> = (0..8).map(|_| r.next_u32()).collect();
    println!("uint32 seed0: {:?}", u);
    let mut r = Mt19937::new(0);
    let ri: Vec<u64> = (0..20).map(|_| r.below(10)).collect();
    println!("randint(0,10,20) seed0: {:?}", ri);
    let mut r = Mt19937::new(1);
    println!("choice(20,5) seed1: {:?}", r.choice_no_replace(20, 5));
    let mut r = Mt19937::new(1);
    println!("permutation(10) seed1: {:?}", r.choice_no_replace(10, 10));
    let mut r = Mt19937::new(1);
    let c = r.choice_no_replace(4500, 4000);
    println!("choice(4500,4000) seed1 first8: {:?}", &c[..8]);
}

const EPS: f64 = 0.005;

/// Build the full log-likelihood matrix logL[r*H + h] for reads that cover >=1 informative site.
/// Reads covering 0 informative positions are dropped (matches em.py loglik_matrix).
pub fn loglik_matrix(rd: &RefData, reads: &[ReadProfile]) -> (Vec<f64>, usize) {
    let h = rd.h;
    let d = (EPS / 3.0).ln() - (1.0 - EPS).ln();
    let cm = (1.0 - EPS).ln();
    let mut rows: Vec<f64> = Vec::new();
    let mut n_reads = 0usize;
    let mut matches = vec![0i32; h];
    for r in reads {
        // count covered informative sites with valid allele
        let mut n = 0i32;
        for x in matches.iter_mut() {
            *x = 0;
        }
        for &(j, c) in r {
            if c < 0 || c > 3 {
                continue;
            }
            n += 1;
            let base = j * h;
            let col = &rd.exp[base..base + h];
            for hh in 0..h {
                if col[hh] == c {
                    matches[hh] += 1;
                }
            }
        }
        if n == 0 {
            continue;
        }
        n_reads += 1;
        let nf = n as f64;
        for hh in 0..h {
            let mm = (n - matches[hh]) as f64;
            rows.push(nf * cm + mm * d);
        }
    }
    (rows, n_reads)
}

/// Core EM on a logL[n*cols] row-major matrix. Uniform log-pi init; 200 iters, tol 1e-6 on total
/// loglik (matches em.py `_em` MLE path). Returns pi (len cols).
pub fn em(logl: &[f64], n: usize, cols: usize, max_iter: usize, tol: f64) -> Vec<f64> {
    if n == 0 {
        return vec![0.0; cols];
    }
    let mut logpi = vec![(1.0f64 / cols as f64).ln(); cols];
    let mut pi = vec![0.0f64; cols];
    let mut prev = f64::NEG_INFINITY;
    let mut gamma_row = vec![0.0f64; cols];
    for _ in 0..max_iter {
        let mut nk = vec![0.0f64; cols];
        let mut ll = 0.0f64;
        for r in 0..n {
            let row = &logl[r * cols..r * cols + cols];
            // joint = logpi + row; softmax
            let mut mx = f64::NEG_INFINITY;
            for h in 0..cols {
                let v = logpi[h] + row[h];
                gamma_row[h] = v;
                if v > mx {
                    mx = v;
                }
            }
            let mut denom = 0.0f64;
            for h in 0..cols {
                let w = (gamma_row[h] - mx).exp();
                gamma_row[h] = w;
                denom += w;
            }
            ll += mx + denom.ln();
            let inv = 1.0 / denom;
            for h in 0..cols {
                nk[h] += gamma_row[h] * inv;
            }
        }
        let nn = n as f64;
        for h in 0..cols {
            pi[h] = nk[h] / nn;
            logpi[h] = pi[h].ln();
        }
        if (ll - prev).abs() < tol * prev.abs().max(1.0) {
            break;
        }
        prev = ll;
    }
    pi
}

/// A contributor component: representative haplogroup index, proportion, supporting read count.
#[derive(Clone, Debug)]
pub struct Comp {
    pub rep: u32,
    pub proportion: f64,
    pub reads: i64,
}

/// Collapse phylogenetically-nested haplogroups (pi>=min_prop) into maximal lineages -> contributors.
/// Mirrors em.py MixtureEM.contributors.
pub fn contributors(rd: &RefData, pi: &[f64], min_prop: f64, n_reads: usize) -> Vec<Comp> {
    let mut kept: Vec<(u32, f64)> = (0..rd.h)
        .filter(|&i| pi[i] >= min_prop)
        .map(|i| (i as u32, pi[i]))
        .collect();
    // sort by -pi (stable; ties keep ascending index like Python's stable sort on original order)
    kept.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    struct Group {
        rep: u32,
        prop: f64,
    }
    let mut groups: Vec<Group> = Vec::new();
    for (hh, p) in kept {
        let mut placed = false;
        for g in groups.iter_mut() {
            if rd.is_ancestor(hh, g.rep) || rd.is_ancestor(g.rep, hh) {
                g.prop += p;
                if rd.lineage[hh as usize].len() > rd.lineage[g.rep as usize].len() {
                    g.rep = hh;
                }
                placed = true;
                break;
            }
        }
        if !placed {
            groups.push(Group { rep: hh, prop: p });
        }
    }
    let tot: f64 = groups.iter().map(|g| g.prop).sum::<f64>().max(1e-300);
    let tot = if tot <= 0.0 { 1.0 } else { tot };
    let mut comps: Vec<Comp> = groups
        .iter()
        .map(|g| Comp {
            rep: g.rep,
            proportion: round4(g.prop / tot),
            reads: (g.prop * n_reads as f64).round() as i64,
        })
        .collect();
    comps.sort_by(|a, b| b.proportion.partial_cmp(&a.proportion).unwrap());
    comps
}

#[inline]
fn round4(x: f64) -> f64 {
    (x * 10000.0).round() / 10000.0
}


/// numpy-compatible MT19937 (legacy RandomState). Seeded via init_by_array([seed]) and drawing
/// bounded integers with numpy's 32-bit masked-rejection algorithm, so `randint(0,N,N)` and
/// `choice(n,cap,replace=False)` reproduce numpy's stream bit-for-bit -> the bootstrap-stability
/// range matches Python em.py exactly (not just within +/-1).
const MT_N: usize = 624;
const MT_M: usize = 397;
const MATRIX_A: u32 = 0x9908b0df;
const UPPER_MASK: u32 = 0x80000000;
const LOWER_MASK: u32 = 0x7fffffff;

struct Mt19937 {
    mt: [u32; MT_N],
    idx: usize,
}

impl Mt19937 {
    fn new(seed: u32) -> Self {
        // numpy RandomState(seed) for a scalar seed in [0, 2^32-1] uses mt19937_seed == init_genrand.
        let mut r = Mt19937 { mt: [0u32; MT_N], idx: MT_N };
        r.init_genrand(seed);
        r
    }

    fn init_genrand(&mut self, s: u32) {
        self.mt[0] = s;
        for i in 1..MT_N {
            let prev = self.mt[i - 1];
            self.mt[i] = (1812433253u32
                .wrapping_mul(prev ^ (prev >> 30)))
                .wrapping_add(i as u32);
        }
        self.idx = MT_N;
    }

    #[allow(dead_code)] // kept for seeds > 2^32-1 (numpy switches to init_by_array); unused for our seeds
    fn init_by_array(&mut self, key: &[u32]) {
        self.init_genrand(19650218);
        let mut i = 1usize;
        let mut j = 0usize;
        let mut k = MT_N.max(key.len());
        while k > 0 {
            let prev = self.mt[i - 1];
            self.mt[i] = (self.mt[i] ^ (prev ^ (prev >> 30)).wrapping_mul(1664525))
                .wrapping_add(key[j])
                .wrapping_add(j as u32);
            i += 1;
            j += 1;
            if i >= MT_N {
                self.mt[0] = self.mt[MT_N - 1];
                i = 1;
            }
            if j >= key.len() {
                j = 0;
            }
            k -= 1;
        }
        let mut k = MT_N - 1;
        while k > 0 {
            let prev = self.mt[i - 1];
            self.mt[i] = (self.mt[i] ^ (prev ^ (prev >> 30)).wrapping_mul(1566083941))
                .wrapping_sub(i as u32);
            i += 1;
            if i >= MT_N {
                self.mt[0] = self.mt[MT_N - 1];
                i = 1;
            }
            k -= 1;
        }
        self.mt[0] = 0x80000000;
        self.idx = MT_N;
    }

    fn generate(&mut self) {
        for i in 0..MT_N {
            let y = (self.mt[i] & UPPER_MASK) | (self.mt[(i + 1) % MT_N] & LOWER_MASK);
            let mut next = self.mt[(i + MT_M) % MT_N] ^ (y >> 1);
            if y & 1 != 0 {
                next ^= MATRIX_A;
            }
            self.mt[i] = next;
        }
        self.idx = 0;
    }

    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.idx >= MT_N {
            self.generate();
        }
        let mut y = self.mt[self.idx];
        self.idx += 1;
        y ^= y >> 11;
        y ^= (y << 7) & 0x9d2c5680;
        y ^= (y << 15) & 0xefc60000;
        y ^= y >> 18;
        y
    }

    /// numpy legacy bounded integer in [0, rng] inclusive via 32-bit masked rejection.
    #[inline]
    fn bounded(&mut self, rng: u32) -> u32 {
        if rng == 0 {
            return 0;
        }
        // mask = smallest 2^k - 1 >= rng
        let mut mask = rng;
        mask |= mask >> 1;
        mask |= mask >> 2;
        mask |= mask >> 4;
        mask |= mask >> 8;
        mask |= mask >> 16;
        loop {
            let v = self.next_u32() & mask;
            if v <= rng {
                return v;
            }
        }
    }

    /// numpy randint(0, n): returns a value in [0, n).
    #[inline]
    fn below(&mut self, n: u64) -> u64 {
        self.bounded((n - 1) as u32) as u64
    }

    /// numpy permutation(n)[..cap]: Fisher-Yates shuffle (legacy shuffle draws j in [0,i]).
    fn choice_no_replace(&mut self, n: usize, cap: usize) -> Vec<usize> {
        let mut arr: Vec<usize> = (0..n).collect();
        let mut i = n - 1;
        while i >= 1 {
            let j = self.bounded(i as u32) as usize;
            arr.swap(i, j);
            i -= 1;
        }
        arr.truncate(cap);
        arr
    }
}

/// The full estimate + bootstrap-stability range (mirrors em.py estimate_range fields we report).
#[derive(Debug)]
pub struct RangeEstimate {
    pub informative_reads: usize,
    pub components: Vec<Comp>,  // raw (deterministic) components
    pub lower_bound_sub: usize, // macro/sub-clade-level contributor lower bound
    pub ci_lo: usize,           // bootstrap 2.5th percentile of nhat
    pub ci_hi: usize,           // bootstrap 97.5th percentile of nhat
}

#[allow(clippy::too_many_arguments)]
pub fn estimate_range(
    rd: &RefData,
    reads: &[ReadProfile],
    min_prop: f64,
    n_boot: usize,
    stab: f64,
    min_reads: f64,
    cap: usize,
    seed: u64,
    boot_iter: usize,
) -> RangeEstimate {
    // drop empty reads
    let mut reads_v: Vec<&ReadProfile> = reads.iter().filter(|r| !r.is_empty()).collect();
    if reads_v.len() > cap {
        // numpy RandomState(1).choice(len, cap, replace=False)
        let mut rng1 = Mt19937::new(1);
        let sub = rng1.choice_no_replace(reads_v.len(), cap);
        reads_v = sub.into_iter().map(|i| reads_v[i]).collect();
    }
    let owned: Vec<ReadProfile> = reads_v.iter().map(|r| (*r).clone()).collect();
    let (logl, n) = loglik_matrix(rd, &owned);
    let h = rd.h;

    // raw estimate (full panel, full iters)
    let raw_pi = em(&logl, n, h, 200, 1e-6);
    let raw_comps = if n > 0 {
        contributors(rd, &raw_pi, min_prop, n)
    } else {
        Vec::new()
    };

    // prune candidate haplogroups for the bootstrap
    let mut cand: Vec<usize> = (0..h).collect();
    if n > 0 {
        let floor_frac = min_reads / n as f64;
        let prune_thresh = (0.25 * floor_frac).min(0.1 * min_prop);
        let c: Vec<usize> = (0..h).filter(|&i| raw_pi[i] >= prune_thresh).collect();
        if !c.is_empty() {
            cand = c;
        }
    }
    let cols = cand.len();
    // logL submatrix over candidate columns
    let mut logl_sub = vec![0.0f64; n * cols];
    for r in 0..n {
        let src = &logl[r * h..r * h + h];
        let dst = &mut logl_sub[r * cols..r * cols + cols];
        for (k, &ci) in cand.iter().enumerate() {
            dst[k] = src[ci];
        }
    }

    // bootstrap stability
    let mut rng = Mt19937::new(seed as u32);
    let mut macro_freq: FxHashMap<String, usize> = FxHashMap::default();
    let mut nhats: Vec<usize> = Vec::with_capacity(n_boot);
    let mut boot_logl = vec![0.0f64; n * cols];
    let mut pi_full = vec![0.0f64; h];
    for _ in 0..n_boot {
        if n == 0 {
            nhats.push(0);
            continue;
        }
        // resample N row indices with replacement
        for r in 0..n {
            let src_row = rng.below(n as u64) as usize;
            boot_logl[r * cols..r * cols + cols]
                .copy_from_slice(&logl_sub[src_row * cols..src_row * cols + cols]);
        }
        let pi_sub = em(&boot_logl, n, cols, boot_iter, 1e-6);
        for v in pi_full.iter_mut() {
            *v = 0.0;
        }
        for (k, &ci) in cand.iter().enumerate() {
            pi_full[ci] = pi_sub[k];
        }
        let comps = contributors(rd, &pi_full, min_prop, n);
        let mut macros: FxHashSet<String> = FxHashSet::default();
        for c in &comps {
            if c.proportion * n as f64 >= min_reads {
                macros.insert(rd.macro_of[c.rep as usize].clone());
            }
        }
        nhats.push(macros.len());
        for m in macros {
            *macro_freq.entry(m).or_insert(0) += 1;
        }
    }

    // stability test: em.py uses `freq / n_boot >= stab` (match that exact float comparison). The
    // count of stable macro-clades is the sub-clade-level contributor lower bound.
    let nb = n_boot as f64;
    let n_stable = macro_freq.values().filter(|&&f| f as f64 / nb >= stab).count();

    let floor = if n > 0 { 1 } else { 0 };
    let lower_bound_sub = n_stable.max(floor);

    // ci95 from the bootstrap distribution of nhats
    let (ci_lo, ci_hi) = if !nhats.is_empty() {
        let mut s = nhats.clone();
        s.sort();
        (
            (percentile_floor(&s, 2.5)).max(floor),
            (percentile_ceil(&s, 97.5)).max(floor),
        )
    } else {
        (floor, floor)
    };

    RangeEstimate {
        informative_reads: n,
        components: raw_comps,
        lower_bound_sub,
        ci_lo,
        ci_hi,
    }
}

/// numpy-style linear-interpolation percentile, then floor (for ci_lo).
fn percentile_floor(sorted: &[usize], q: f64) -> usize {
    percentile(sorted, q).floor() as usize
}
fn percentile_ceil(sorted: &[usize], q: f64) -> usize {
    percentile(sorted, q).ceil() as usize
}
fn percentile(sorted: &[usize], q: f64) -> f64 {
    let n = sorted.len();
    if n == 0 {
        return 0.0;
    }
    if n == 1 {
        return sorted[0] as f64;
    }
    let rank = (q / 100.0) * (n - 1) as f64;
    let lo = rank.floor() as usize;
    let hi = rank.ceil() as usize;
    let frac = rank - lo as f64;
    sorted[lo] as f64 + (sorted[hi] as f64 - sorted[lo] as f64) * frac
}