mtnoc-rs 0.1.0

//! Custom seed-anchored short-read aligner, FUSED with per-base allele extraction, built on the
//! `triple_accel` SIMD edit-distance kernels. Validated against the strobealign chrM BAM.
//!
//! Pipeline per read:
//!   1. Seed: find the first forward 21-mer (from kmer::kmers_fwd) that hits the chrM seed index,
//!      in both the read and its reverse-complement orientation. The hit fixes an anchor mapping
//!      read-base `off` -> ref-pos `rpos`, hence an ungapped ref start = rpos - off.
//!   2. Ungapped placement: slice the ref to the read's span at that start and count mismatches with
//!      triple_accel::hamming (SIMD). If mismatch-rate over the overlap <= 0.05 -> accept ungapped,
//!      with an exact per-base map (read base i -> ref pos start+i) => fused alignment + alleles.
//!   3. Fallback (indels / high mismatch): triple_accel::levenshtein_search of the read against a
//!      banded ref window around the anchor -> best edit distance k and its ref start/end.
//!   Best of the two orientations wins. A read "maps to chrM" if best identity >= min_identity.

use crate::kmer::{self, K};
use anyhow::{Context, Result};
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::time::Instant;
use triple_accel::{hamming, levenshtein_search};

/// chrM 21-mer -> first ref position seed index (forward strand only; reads are oriented instead).
/// Also holds NUMT decoy contigs + their seed index so a read that aligns BETTER to a decoy than to
/// chrM can be rejected (competitive best-contig assignment, like strobealign to the compact ref).
pub struct RefIndex {
    seq: Vec<u8>,
    idx: FxHashMap<u64, u32>,
    decoys: Vec<Vec<u8>>,
    didx: FxHashMap<u64, (u32, u32)>, // decoy 21-mer -> (contig, pos)
}

impl RefIndex {
    /// Full compact-ref index: chrM (`chrm_contig`) for alignment + all other contigs as NUMT decoys
    /// for competitive rejection.
    pub fn build_multi(path: &str, chrm_contig: &str) -> Result<Self> {
        let (mut seq, mut decoys) = (Vec::new(), Vec::new());
        let mut cur: Vec<u8> = Vec::new();
        let mut cur_is_chrm = false;
        let f = BufReader::new(File::open(path).with_context(|| format!("open {path}"))?);
        let flush = |cur: &mut Vec<u8>, is_chrm: bool, seq: &mut Vec<u8>, decoys: &mut Vec<Vec<u8>>| {
            if cur.is_empty() {
                return;
            }
            if is_chrm {
                *seq = std::mem::take(cur);
            } else {
                decoys.push(std::mem::take(cur));
            }
        };
        for line in f.lines() {
            let line = line?;
            if let Some(h) = line.strip_prefix('>') {
                flush(&mut cur, cur_is_chrm, &mut seq, &mut decoys);
                cur_is_chrm = h.split_whitespace().next() == Some(chrm_contig);
            } else {
                cur.extend_from_slice(line.trim().as_bytes());
            }
        }
        flush(&mut cur, cur_is_chrm, &mut seq, &mut decoys);
        anyhow::ensure!(!seq.is_empty(), "chrM contig {chrm_contig} not found in {path}");

        let mut idx = FxHashMap::default();
        for (pos, code) in kmer::kmers_fwd(&seq) {
            idx.entry(code).or_insert(pos as u32);
        }
        let mut didx: FxHashMap<u64, (u32, u32)> = FxHashMap::default();
        for (ci, d) in decoys.iter().enumerate() {
            for (pos, code) in kmer::kmers_fwd(d) {
                didx.entry(code).or_insert((ci as u32, pos as u32));
            }
        }
        Ok(RefIndex { seq, idx, decoys, didx })
    }

    /// Best ungapped identity of `read` (either orientation) against any NUMT decoy it seeds to.
    /// Used to reject reads better explained by a NUMT than by chrM.
    fn best_decoy_ident(&self, read: &[u8]) -> f64 {
        if self.decoys.is_empty() {
            return 0.0;
        }
        let rc = kmer::revcomp_seq(read);
        let mut best = 0.0f64;
        for q in [read, rc.as_slice()] {
            let l = q.len();
            if l < K {
                continue;
            }
            if let Some((off, cid, rp)) = kmer::kmers_fwd(q)
                .into_iter()
                .find_map(|(off, code)| self.didx.get(&code).map(|&(c, p)| (off, c as usize, p as usize)))
            {
                let dseq = &self.decoys[cid];
                let dl = dseq.len();
                let ref_start = rp as i64 - off as i64;
                let read_lo = (-ref_start).max(0) as usize;
                let ref_lo = ref_start.max(0) as usize;
                let ref_hi = ((ref_start + l as i64) as usize).min(dl);
                if ref_hi > ref_lo {
                    let ov = ref_hi - ref_lo;
                    let mm = hamming(&q[read_lo..read_lo + ov], &dseq[ref_lo..ref_hi]);
                    let ident = 1.0 - (mm as usize + (l - ov)) as f64 / l as f64;
                    if ident > best {
                        best = ident;
                    }
                }
            }
        }
        best
    }
}

/// How an accepted alignment was produced.
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Mode {
    Ungapped,
    Levenshtein,
}

/// One accepted alignment: leftmost ref start (0-based), identity, mode.
#[derive(Clone, Copy, Debug)]
pub struct Aln {
    pub ref_start0: usize, // 0-based leftmost ref coordinate of the alignment
    pub ident: f64,
    pub mode: Mode,
}

const BAND: usize = 25; // ref-window slack for the indel fallback

/// Align one oriented query (already fwd or rc). Returns the best placement if a seed anchors.
fn align_oriented(index: &RefIndex, query: &[u8]) -> Option<Aln> {
    let reflen = index.seq.len();
    let l = query.len();
    if l < K {
        return None;
    }
    // 1. anchor: first forward 21-mer of the query that hits the chrM index.
    let (off, rpos) = kmer::kmers_fwd(query)
        .into_iter()
        .find_map(|(off, code)| index.idx.get(&code).map(|&rp| (off, rp as usize)))?;

    // 2. ungapped: read base 0 -> ref (rpos - off).
    let ref_start = rpos as i64 - off as i64;
    // overlap of [ref_start, ref_start+l) with [0, reflen)
    let read_lo = (-ref_start).max(0) as usize;
    let ref_lo = ref_start.max(0) as usize;
    let ref_hi = ((ref_start + l as i64) as usize).min(reflen);
    if ref_hi > ref_lo {
        let ov = ref_hi - ref_lo; // overlap length
        let read_hi = read_lo + ov;
        let mm = hamming(&query[read_lo..read_hi], &index.seq[ref_lo..ref_hi]);
        // overhang bases (read hanging off the ref end) count as non-matching for identity.
        let total_edits = mm as usize + (l - ov);
        let ident = 1.0 - total_edits as f64 / l as f64;
        // accept ungapped when the aligned span is nearly full and clean (<=5% mismatch over overlap).
        if ov >= (l * 9) / 10 && (mm as f64) <= 0.05 * ov as f64 {
            return Some(Aln {
                ref_start0: ref_lo,
                ident,
                mode: Mode::Ungapped,
            });
        }
    }

    // 3. fallback: banded Levenshtein search (handles indels + higher divergence).
    let win_lo = (ref_start - BAND as i64).max(0) as usize;
    let win_hi = ((ref_start + l as i64 + BAND as i64) as usize).min(reflen);
    if win_hi <= win_lo || win_hi - win_lo < l / 2 {
        return None;
    }
    let best = levenshtein_search(query, &index.seq[win_lo..win_hi]).min_by_key(|m| m.k)?;
    let ident = 1.0 - best.k as f64 / l as f64;
    Some(Aln {
        ref_start0: win_lo + best.start,
        ident,
        mode: Mode::Levenshtein,
    })
}

/// Align a read in both orientations; keep the higher-identity placement. Reject reads that align
/// BETTER to a NUMT decoy than to chrM (competitive best-contig assignment).
pub fn align_read(index: &RefIndex, read: &[u8]) -> Option<Aln> {
    let fwd = align_oriented(index, read);
    let rc = kmer::revcomp_seq(read);
    let rev = align_oriented(index, &rc);
    let best = match (fwd, rev) {
        (Some(a), Some(b)) => Some(if a.ident >= b.ident { a } else { b }),
        (Some(a), None) => Some(a),
        (None, Some(b)) => Some(b),
        (None, None) => None,
    }?;
    if index.best_decoy_ident(read) > best.ident {
        return None; // better explained by a NUMT
    }
    Some(best)
}

// ---------------------------------------------------------------------------
// loading
// ---------------------------------------------------------------------------

/// Align a read in both orientations and, for the winning placement, also emit the per-base allele
/// map (1-based ref position, uppercase base) along the accepted alignment diagonal. For ungapped
/// placements the map is exact; for the (rare) Levenshtein fallback it approximates the diagonal
/// (ignores the indel offset) — acceptable since informative positions are sparse.
pub fn align_with_alleles(index: &RefIndex, read: &[u8], qual: &[u8]) -> Option<(Aln, Vec<(i64, u8)>)> {
    let fwd = extract_oriented(index, read, qual);
    // reverse-complement the sequence AND reverse the quality string for the rc orientation
    let rc = kmer::revcomp_seq(read);
    let rq: Vec<u8> = qual.iter().rev().copied().collect();
    let rev = extract_oriented(index, &rc, &rq);
    let best = match (fwd, rev) {
        (Some(a), Some(b)) => Some(if a.0.ident >= b.0.ident { a } else { b }),
        (Some(a), None) => Some(a),
        (None, Some(b)) => Some(b),
        (None, None) => None,
    }?;
    if index.best_decoy_ident(read) > best.0.ident {
        return None; // better explained by a NUMT decoy
    }
    Some(best)
}

const MIN_BQ: u8 = 20 + 33; // Phred>=20 (mpileup -Q 20); qual bytes are Phred+33 ASCII

/// Like `align_oriented` but also returns the per-base (1-based ref pos, base) list for the placement,
/// keeping only bases with quality >= MIN_BQ (matching the Python mpileup -Q 20 allele extraction).
fn extract_oriented(index: &RefIndex, query: &[u8], qual: &[u8]) -> Option<(Aln, Vec<(i64, u8)>)> {
    let reflen = index.seq.len();
    let l = query.len();
    if l < K {
        return None;
    }
    let (off, rpos) = kmer::kmers_fwd(query)
        .into_iter()
        .find_map(|(off, code)| index.idx.get(&code).map(|&rp| (off, rp as usize)))?;

    let ref_start = rpos as i64 - off as i64;
    let read_lo = (-ref_start).max(0) as usize;
    let ref_lo = ref_start.max(0) as usize;
    let ref_hi = ((ref_start + l as i64) as usize).min(reflen);
    if ref_hi > ref_lo {
        let ov = ref_hi - ref_lo;
        let read_hi = read_lo + ov;
        let mm = hamming(&query[read_lo..read_hi], &index.seq[ref_lo..ref_hi]);
        let total_edits = mm as usize + (l - ov);
        let ident = 1.0 - total_edits as f64 / l as f64;
        if ov >= (l * 9) / 10 && (mm as f64) <= 0.05 * ov as f64 {
            // exact per-base map: query[read_lo+t] -> ref (ref_lo + t), gated by base quality
            let mut alleles = Vec::with_capacity(ov);
            for t in 0..ov {
                if qual.get(read_lo + t).map_or(false, |&q| q >= MIN_BQ) {
                    alleles.push(((ref_lo + t) as i64 + 1, query[read_lo + t].to_ascii_uppercase()));
                }
            }
            return Some((
                Aln { ref_start0: ref_lo, ident, mode: Mode::Ungapped },
                alleles,
            ));
        }
    }

    // Levenshtein fallback
    let win_lo = (ref_start - BAND as i64).max(0) as usize;
    let win_hi = ((ref_start + l as i64 + BAND as i64) as usize).min(reflen);
    if win_hi <= win_lo || win_hi - win_lo < l / 2 {
        return None;
    }
    let best = levenshtein_search(query, &index.seq[win_lo..win_hi]).min_by_key(|m| m.k)?;
    let ident = 1.0 - best.k as f64 / l as f64;
    let ref_start0 = win_lo + best.start;
    // approximate diagonal: query[i] -> ref (ref_start0 + i), gated by base quality
    let span = l.min(reflen.saturating_sub(ref_start0));
    let mut alleles = Vec::with_capacity(span);
    for i in 0..span {
        if qual.get(i).map_or(false, |&q| q >= MIN_BQ) {
            alleles.push(((ref_start0 + i) as i64 + 1, query[i].to_ascii_uppercase()));
        }
    }
    Some((
        Aln { ref_start0, ident, mode: Mode::Levenshtein },
        alleles,
    ))
}

/// One candidate read: QNAME (commas->underscores, to match the BAM), mate (1/2), sequence + quality.
pub struct Read {
    pub name: String,
    pub mate: u8,
    pub seq: Vec<u8>,
    pub qual: Vec<u8>, // raw Phred+33 ASCII, same length as seq
}

pub fn load_reads(path: &str, mate: u8) -> Result<Vec<Read>> {
    let mut r = BufReader::new(File::open(path).with_context(|| format!("open {path}"))?);
    let mut out = Vec::new();
    let (mut name, mut seq, mut i) = (String::new(), Vec::new(), 0usize);
    let mut line = String::new();
    loop {
        line.clear();
        if r.read_line(&mut line)? == 0 {
            break;
        }
        match i % 4 {
            0 => {
                // '@QNAME ...'; QNAME is up to first whitespace; sanitize commas -> underscores.
                name = line[1..].split_whitespace().next().unwrap_or("").replace(',', "_");
            }
            1 => seq = line.trim_end().as_bytes().to_vec(),
            3 => out.push(Read {
                name: name.clone(),
                mate,
                seq: std::mem::take(&mut seq),
                qual: line.trim_end().as_bytes().to_vec(),
            }),
            _ => {}
        }
        i += 1;
    }
    Ok(out)
}

/// Run the aligner over the candidate fastqs and report agreement vs the BAM QNAME set.
/// `truth`: set of QNAMEs mapped to chrM in the BAM (sanitized). `truth_pos`: (qname,mate)->1-based POS.
pub fn run_validate(
    reference: &str,
    contig: &str,
    c1: &str,
    c2: &str,
    min_identity: f64,
    truth: &std::collections::HashSet<String>,
    truth_pos: &FxHashMap<(String, u8), i64>,
) -> Result<()> {
    let index = RefIndex::build_multi(reference, contig)?;

    let mut reads = load_reads(c1, 1)?;
    reads.extend(load_reads(c2, 2)?);
    eprintln!(
        "[validate] chrM {} bp ({} 21-mers) + {} NUMT decoys, candidate reads {} (mates)",
        index.seq.len(),
        index.idx.len(),
        index.decoys.len(),
        reads.len()
    );

    crate::mem_reset_peak();
    let t = Instant::now();
    // fused alignment + allele extraction, parallel over independent reads.
    let results: Vec<(usize, Option<Aln>)> = reads
        .par_iter()
        .enumerate()
        .map(|(i, rd)| (i, align_read(&index, &rd.seq)))
        .collect();
    let dt = t.elapsed().as_secs_f64();
    let mem = crate::mem_peak_delta();

    // fragment-level mapped set: a QNAME maps if either mate aligns at identity >= min_identity.
    use std::collections::HashSet;
    let mut mapped: HashSet<String> = HashSet::new();
    let mut n_mate_mapped = 0usize;
    let mut n_ungapped = 0usize;
    let mut n_leven = 0usize;
    // (qname,mate)->my 1-based pos, for the position spot-check
    let mut my_pos: FxHashMap<(String, u8), i64> = FxHashMap::default();
    for (i, res) in &results {
        if let Some(a) = res {
            if a.ident >= min_identity {
                n_mate_mapped += 1;
                match a.mode {
                    Mode::Ungapped => n_ungapped += 1,
                    Mode::Levenshtein => n_leven += 1,
                }
                let rd = &reads[*i];
                mapped.insert(rd.name.clone());
                my_pos.insert((rd.name.clone(), rd.mate), a.ref_start0 as i64 + 1);
            }
        }
    }

    // ---- read-name agreement (fragment level) ----
    let inter = mapped.intersection(truth).count();
    let union = mapped.len() + truth.len() - inter;
    let recall = 100.0 * inter as f64 / truth.len().max(1) as f64;
    let precision = 100.0 * inter as f64 / mapped.len().max(1) as f64;
    let jaccard = 100.0 * inter as f64 / union.max(1) as f64;

    println!("=== custom triple_accel seed-anchored aligner vs strobealign chrM BAM ===");
    println!("reads (mates) processed : {}", reads.len());
    println!("mate alignments accepted: {n_mate_mapped} (ungapped {n_ungapped}, levenshtein {n_leven})");
    println!("fragments mapped (mine) : {}", mapped.len());
    println!("fragments in BAM (truth): {}", truth.len());
    println!("intersection            : {inter}");
    println!("recall vs BAM           : {recall:.2}%  (BAM fragments recovered)");
    println!("precision vs BAM        : {precision:.2}%  (mine also in BAM; NUMTs inflate)");
    println!("Jaccard (inter/union)   : {jaccard:.2}%");
    println!("wall time (align)       : {dt:.3}s   ({:.0} mate-reads/s)", reads.len() as f64 / dt);
    println!("peak add'l memory       : {:.1} MB", mem as f64 / 1e6);

    // ---- position spot-check on shared (qname,mate) pairs ----
    let mut shared: Vec<((String, u8), i64, i64)> = Vec::new();
    for (key, &mp) in &my_pos {
        if let Some(&bp) = truth_pos.get(key) {
            shared.push((key.clone(), mp, bp));
        }
    }
    shared.sort();
    let sample_n = 30.min(shared.len());
    let within2 = shared.iter().filter(|(_, mp, bp)| (mp - bp).abs() <= 2).count();
    let within5 = shared.iter().filter(|(_, mp, bp)| (mp - bp).abs() <= 5).count();
    println!("\n--- position spot-check ((qname,mate) present in both, with a BAM POS) ---");
    println!(
        "shared placements: {}   |Δpos|<=2bp: {} ({:.1}%)   <=5bp: {} ({:.1}%)",
        shared.len(),
        within2,
        100.0 * within2 as f64 / shared.len().max(1) as f64,
        within5,
        100.0 * within5 as f64 / shared.len().max(1) as f64
    );
    println!("sample (mine_POS vs BAM_POS, mate):");
    for ((name, mate), mp, bp) in shared.iter().take(sample_n) {
        let flag = if (mp - bp).abs() <= 2 { "ok" } else { "DIFF" };
        let short = if name.len() > 40 { &name[name.len() - 40..] } else { name.as_str() };
        println!("  {short:>40} m{mate}  mine={mp:>6}  bam={bp:>6}  {flag}");
    }
    Ok(())
}

#[allow(dead_code)]
const _K: usize = K;