mtnoc-rs 0.1.1

//! mtnoc-rs — pure-Rust reimplementation of mtNoC.
//! `fast` (100% pure Rust): raw FASTQ pair -> embedded-rCRS k-mer prefilter (R1/R2 decompressed+matched
//! on separate threads, a consumer keeps a pair when EITHER mate shares a canonical 21-mer) ->
//! triple_accel seed-anchored aligner -> mixture-EM. `accurate`: strobealign + samtools -> same EM.

mod em;
mod kmer;
mod myaln;

use anyhow::{Context, Result};
// Production uses the system allocator; these mem-tracking hooks are no-ops (they existed only for the
// one-off aligner-crate memory benchmark, whose atomic accounting slowed every allocation).
pub fn mem_reset_peak() {}
pub fn mem_peak_delta() -> usize {
    0
}
use clap::{Args, Parser, Subcommand};
use flate2::read::MultiGzDecoder;
use kmer::KmerSet;
use needletail::parse_fastx_reader;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Write};
use std::process::{Command, Stdio};
use std::sync::mpsc::sync_channel;
use std::sync::Arc;
use std::thread;

// Reference tables baked into the binary (gzip = pure-Rust flate2/zlib-rs, no C deps). These are
// build-independent data, so embedding keeps `fast`/`accurate` self-contained: the user only supplies
// the alignment FASTA via --ref. Regenerate with `pigz -9` from data/{mtref.txt,rcrs_wrap.fa}.
const EMB_MTREF_GZ: &[u8] = include_bytes!("embedded/mtref.txt.gz");
const EMB_RCRS_GZ: &[u8] = include_bytes!("embedded/rcrs_wrap.fa.gz");

/// Decompress an embedded gzip asset fully into memory.
fn gunzip(bytes: &'static [u8]) -> Result<Vec<u8>> {
    let mut out = Vec::new();
    MultiGzDecoder::new(bytes)
        .read_to_end(&mut out)
        .context("decompress embedded reference")?;
    Ok(out)
}

/// The embedded EM reference dump (dump_ref_for_rust.py output), decompressed and parsed.
fn embedded_refdata() -> Result<em::RefData> {
    let raw = gunzip(EMB_MTREF_GZ)?;
    em::RefData::load_reader(BufReader::new(&raw[..]))
}

/// The embedded circular-wrapped rCRS k-mer prefilter reference sequence (concatenated, no header).
fn embedded_kmer_ref() -> Result<Vec<u8>> {
    let raw = gunzip(EMB_RCRS_GZ)?;
    let mut seq = Vec::new();
    for line in raw.split(|&b| b == b'\n') {
        if line.first() == Some(&b'>') {
            if !seq.is_empty() {
                break;
            }
        } else {
            seq.extend(line.iter().filter(|b| !b.is_ascii_whitespace()));
        }
    }
    Ok(seq)
}

#[derive(Parser)]
#[command(name = "mtnoc", version, about = "mtNoC — de-identified mtDNA contributor lower bound")]
struct Cli {
    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Subcommand)]
enum Cmd {
    /// FAST (100% pure Rust): raw FASTQ pair(s) -> k-mer prefilter -> triple_accel seed-anchored
    /// aligner -> mixture-EM -> contributor lower bound. No external tools.
    Fast(EstimateArgs),
    /// ACCURATE: raw FASTQ pair(s) -> strobealign -> samtools filter (MAPQ>=20 & NM/qlen<=0.027) +
    /// mpileup allele extraction -> the SAME mixture-EM. Replicates the reference cohort exactly.
    /// Shells out to `strobealign` + `samtools` (pass `--bam` to skip strobealign).
    Accurate(EstimateArgs),

    /// [dev] Align candidate reads and report recovery vs a truth BAM's read names/positions.
    #[command(hide = true)]
    Validate {
        #[arg(long)]
        reference: String,
        #[arg(long, default_value = "NC_012920.1")]
        contig: String,
        #[arg(long)]
        c1: String,
        #[arg(long)]
        c2: String,
        #[arg(long, default_value_t = 0.90)]
        min_identity: f64,
        #[arg(long)]
        truth_names: String,
        #[arg(long)]
        truth_pos: String,
    },
    /// [dev] RNG self-check (compare MT19937 output against numpy RandomState).
    #[command(hide = true)]
    Rngtest,
    /// [dev] EM-port validation: run the EM directly on a per-fragment allele-profile dump.
    /// Profiles file: one fragment per line, `qname<TAB>pos:allele,pos:allele,...` (1-based pos).
    #[command(hide = true)]
    Emtest {
        #[arg(long)]
        profiles: String,
        #[arg(long, default_value_t = 0.02)]
        min_prop: f64,
        #[arg(long, default_value_t = 30)]
        n_boot: usize,
        #[arg(long, default_value_t = 0)]
        seed: u64,
    },
}

/// Shared arguments for `fast` and `accurate`. One or more samples in, one TSV row each out.
///
/// Inputs: give paired files explicitly with `--in1/--in2` (comma-separate multiple samples), or use
/// the shorthand `--in` with sample stems (the tool appends `_R1`/`_R2` + a FASTQ extension). Sample
/// names are derived from the R1 filename (strip the `_R1...` suffix). `--out1/--out2` (or `--out`)
/// optionally writes the reads that aligned to chrM; omit to only emit the estimate.
#[derive(Args)]
struct EstimateArgs {
    /// Paired input mate 1 FASTQ(.gz). Comma-separate for multiple samples. Use with --in2.
    #[arg(long, value_delimiter = ',', conflicts_with = "in_stem")]
    in1: Vec<String>,
    /// Paired input mate 2 FASTQ(.gz). Comma-separate for multiple samples. Use with --in1.
    #[arg(long, value_delimiter = ',', conflicts_with = "in_stem")]
    in2: Vec<String>,
    /// Shorthand: sample stem(s) without the `_R1`/`_R2` suffix; the tool resolves both mates.
    /// Comma-separate for multiple samples. Mutually exclusive with --in1/--in2.
    #[arg(long = "in", value_delimiter = ',')]
    in_stem: Vec<String>,

    /// Optional: write chrM-aligned (candidate) mate-1 reads here (parallel to inputs).
    #[arg(long, value_delimiter = ',')]
    out1: Vec<String>,
    /// Optional: write chrM-aligned (candidate) mate-2 reads here (parallel to inputs).
    #[arg(long, value_delimiter = ',')]
    out2: Vec<String>,
    /// Optional shorthand for --out1/--out2: output stem(s); the tool appends `_R1`/`_R2`.fq.
    #[arg(long = "out", value_delimiter = ',')]
    out_stem: Vec<String>,

    /// Alignment reference FASTA (compact chrM+NUMT-decoy, shipped as data/mtNoC_compact.fa).
    #[arg(long)]
    r#ref: String,
    /// Reference contig name (chrM) in --ref.
    #[arg(long, default_value = "NC_012920.1")]
    contig: String,

    /// [accurate] pre-aligned, sorted+indexed chrM BAM(s) (.bai alongside); skips strobealign.
    #[arg(long, value_delimiter = ',')]
    bam: Vec<String>,
    /// [accurate] strobealign binary (looked up on PATH by default).
    #[arg(long, default_value = "strobealign")]
    strobealign: String,
    /// [accurate] samtools binary (looked up on PATH by default).
    #[arg(long, default_value = "samtools")]
    samtools: String,

    /// Minimum per-read alignment identity to accept (fast mode).
    #[arg(long, default_value_t = 0.90)]
    min_identity: f64,
    /// Drop EM components below this mixture proportion.
    #[arg(long, default_value_t = 0.02)]
    min_prop: f64,
    /// Bootstrap resamples for the stability range.
    #[arg(long, default_value_t = 30)]
    n_boot: usize,
    /// RNG seed for the bootstrap.
    #[arg(long, default_value_t = 0)]
    seed: u64,
    /// Worker threads (0 = all cores).
    #[arg(long, default_value_t = 0)]
    thread: usize,
    /// Print a column-name header line before the results.
    #[arg(long)]
    header: bool,
}

fn base_code(b: u8) -> i8 {
    match b {
        b'A' | b'a' => 0,
        b'C' | b'c' => 1,
        b'G' | b'g' => 2,
        b'T' | b't' => 3,
        _ => -1,
    }
}

/// Convert (1-based pos, base) observations into (info-index, code) EM entries (drop non-info / non-ACGT).
fn to_read_profile(rd: &em::RefData, obs: &[(i64, u8)]) -> em::ReadProfile {
    let mut prof = em::ReadProfile::new();
    for &(pos, base) in obs {
        if let Some(&j) = rd.info_idx.get(&pos) {
            let c = base_code(base);
            if c >= 0 {
                prof.push((j, c));
            }
        }
    }
    prof
}

/// Output columns (TSV). `--header` prints these names; values are one row per sample.
const OUT_HEADER: &str =
    "sample\tinformative_reads\tcontributors_ci95\tcomposition";

/// Format one result row. `contributors_ci95` = `ci95_low;min_contributors_subclade;ci95_high`;
/// `composition` = `haplogroup:proportion:reads` components joined by `;`.
fn format_row(sample: &str, rd: &em::RefData, est: &em::RangeEstimate) -> String {
    let comp: Vec<String> = est
        .components
        .iter()
        .map(|c| format!("{}:{:.4}:{}r", rd.haplo[c.rep as usize], c.proportion, c.reads))
        .collect();
    let ci = format!("{};{};{}", est.ci_lo, est.lower_bound_sub, est.ci_hi);
    format!("{sample}\t{}\t{ci}\t{}", est.informative_reads, comp.join(";"))
}

/// FAST mode for one sample: raw FASTQ pair -> embedded-rCRS k-mer prefilter -> triple_accel aligner
/// (vs --ref) -> per-read alleles -> mixture-EM. Fully in-memory (no intermediate FASTQs); if
/// `out1`/`out2` are given, the chrM-candidate reads are also written there.
#[allow(clippy::too_many_arguments)]
fn estimate_fast(
    rd: &em::RefData,
    kset: &Arc<KmerSet>,
    index: &myaln::RefIndex,
    r1: &str,
    r2: &str,
    out1: Option<&str>,
    out2: Option<&str>,
    min_identity: f64,
    min_prop: f64,
    n_boot: usize,
    seed: u64,
) -> Result<em::RangeEstimate> {
    let reads = prefilter(kset, r1, r2, out1, out2, 1)?;
    eprintln!("[fast] candidate mate-reads: {}", reads.len());

    use rayon::prelude::*;
    let per_mate: Vec<Option<(String, Vec<(i64, u8)>)>> = reads
        .par_iter()
        .map(|rd_read| {
            myaln::align_with_alleles(index, &rd_read.seq, &rd_read.qual).and_then(|(a, alleles)| {
                if a.ident >= min_identity {
                    Some((rd_read.name.clone(), alleles))
                } else {
                    None
                }
            })
        })
        .collect();

    // group by fragment QNAME (both mates combined), like mpileup --output-QNAME per read name
    use rustc_hash::FxHashMap;
    let mut frags: FxHashMap<String, em::ReadProfile> = FxHashMap::default();
    let mut n_mapped = 0usize;
    for (name, obs) in per_mate.into_iter().flatten() {
        n_mapped += 1;
        let prof = to_read_profile(rd, &obs);
        frags.entry(name).or_default().extend(prof);
    }
    eprintln!("[fast] mate alignments accepted: {n_mapped}; fragments: {}", frags.len());

    let reads_v: Vec<em::ReadProfile> = frags.into_values().collect();
    Ok(em::estimate_range(rd, &reads_v, min_prop, n_boot, 0.8, 8.0, 4000, seed, 200))
}

fn run_emtest(
    profiles: &str,
    min_prop: f64,
    n_boot: usize,
    seed: u64,
) -> Result<()> {
    let rd = embedded_refdata()?;
    let f = BufReader::new(File::open(profiles)?);
    let mut reads_v: Vec<em::ReadProfile> = Vec::new();
    for line in f.lines() {
        let line = line?;
        let mut it = line.splitn(2, '\t');
        let _qname = it.next().unwrap_or("");
        let rest = it.next().unwrap_or("");
        if rest.is_empty() {
            continue;
        }
        let mut obs: Vec<(i64, u8)> = Vec::new();
        for tok in rest.split(',') {
            let mut pa = tok.split(':');
            if let (Some(p), Some(a)) = (pa.next(), pa.next()) {
                if let Ok(pos) = p.parse::<i64>() {
                    let base = a.as_bytes().first().copied().unwrap_or(b'N');
                    obs.push((pos, base));
                }
            }
        }
        reads_v.push(to_read_profile(&rd, &obs));
    }
    eprintln!("[emtest] fragments loaded: {}", reads_v.len());
    let est = em::estimate_range(&rd, &reads_v, min_prop, n_boot, 0.8, 8.0, 4000, seed, 200);
    println!("{OUT_HEADER}");
    println!("{}", format_row("emtest", &rd, &est));
    Ok(())
}

// ---------------------------------------------------------------------------
// ACCURATE MODE: strobealign chrM BAM -> samtools filter + mpileup -> Rust EM.
//
// Faithfully replicates the Python cohort path (src/mtnoc/{filters,readmodel}.py invoked by
// em.py's __main__): (1) `samtools view -e '[NM]/qlen<=0.027 && mapq>=20'` with the QNAME comma
// sanitiser, (2) `samtools mpileup -q 20 -Q 20 --output-QNAME` parsed with readmodel's
// _parse_pileup_bases logic (restricted to informative positions = sites), grouped per read in
// first-seen order (== Python dict insertion order, which the cap subsample & bootstrap depend on),
// (3) the SAME em::estimate_range. Shelling to samtools is the explicit non-pure, max-accuracy path;
// it is a runtime subprocess only (no C crates enter the build -> fast mode stays 100% pure Rust).
// ---------------------------------------------------------------------------

/// Port of readmodel.py `_parse_pileup_bases`: tokenise an mpileup base string into exactly one
/// allele per covering read (so it zips 1:1 with the `--output-QNAME` read list). Emits upper-case
/// A/C/G/T (incl. `.`/`,` resolved to `ref_base`), `N` for ambiguous, `*` for a deletion.
fn parse_pileup_bases(bases: &[u8], ref_base: u8) -> Vec<u8> {
    let ref_base = ref_base.to_ascii_uppercase();
    let mut out: Vec<u8> = Vec::new();
    let n = bases.len();
    let mut i = 0usize;
    while i < n {
        let ch = bases[i];
        match ch {
            b'^' => i += 2, // segment start: skip the mapping-quality char
            b'$' => i += 1, // segment end marker
            b'+' | b'-' => {
                // indel: read the length then skip that many chars
                let mut j = i + 1;
                let mut num = 0usize;
                let mut has = false;
                while j < n && bases[j].is_ascii_digit() {
                    num = num * 10 + (bases[j] - b'0') as usize;
                    j += 1;
                    has = true;
                }
                i = j + if has { num } else { 0 };
            }
            b'.' | b',' => {
                out.push(ref_base);
                i += 1;
            }
            b'A' | b'C' | b'G' | b'T' | b'a' | b'c' | b'g' | b't' => {
                out.push(ch.to_ascii_uppercase());
                i += 1;
            }
            b'N' | b'n' => {
                out.push(b'N');
                i += 1;
            }
            b'*' | b'#' => {
                out.push(b'*');
                i += 1;
            }
            _ => i += 1, // '>' '<' ref-skip, stray chars
        }
    }
    out
}

/// Replicate filters.py `filter_bam` fast path: keep reads with MAPQ>=20 AND NM/qlen<=0.027 in
/// `contig`, sanitise commas in QNAME (so mpileup's comma-joined QNAME column stays parseable), and
/// re-encode + index. Uses the identical samtools filter-expression + awk pipeline the cohort used.
fn filter_bam_accurate(samtools: &str, in_bam: &str, out_bam: &str, contig: &str) -> Result<()> {
    // exact expr string from filters.py (rate=0.027, mapq_min=20)
    let pipeline = format!(
        "{s} view -h -e '[NM]/qlen <= 0.027 && mapq >= 20' \"{i}\" {c} | \
         awk 'BEGIN{{OFS=\"\\t\"}} /^@/{{print;next}} {{gsub(/,/,\"_\",$1); print}}' | \
         {s} view -b -o \"{o}\" -",
        s = samtools,
        i = in_bam,
        c = contig,
        o = out_bam
    );
    let status = Command::new("bash")
        .arg("-c")
        .arg(&pipeline)
        .status()
        .context("spawn samtools filter pipeline")?;
    anyhow::ensure!(status.success(), "samtools filter pipeline failed");
    let status = Command::new(samtools)
        .args(["index", out_bam])
        .status()
        .context("spawn samtools index")?;
    anyhow::ensure!(status.success(), "samtools index failed");
    Ok(())
}

/// Replicate readmodel.py `per_read_profiles` (sites=informative, ref_alleles=True) over
/// `samtools mpileup -f <ref> -r <contig> -q 20 -Q 20 --output-QNAME <bam>`. Returns per-fragment
/// allele profiles as em::ReadProfile in FIRST-SEEN read order (== Python dict order). Non-ACGT
/// alleles (N/`*`) are kept as sentinel entries `(j, -1)` so a read observed ONLY at N/`*` info
/// sites still counts as a non-empty fragment (matching Python, where the 4000-read cap subsample
/// is taken over ALL profiles before the log-likelihood build drops uninformative ones).
fn pileup_profiles(
    rd: &em::RefData,
    samtools: &str,
    reference: &str,
    contig: &str,
    bam: &str,
) -> Result<Vec<em::ReadProfile>> {
    let mut child = Command::new(samtools)
        .args([
            "mpileup",
            "-f",
            reference,
            "-r",
            contig,
            "-q",
            "20",
            "-Q",
            "20",
            "--output-QNAME",
            bam,
        ])
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
        .context("spawn samtools mpileup")?;
    let stdout = child.stdout.take().context("mpileup stdout")?;
    let reader = BufReader::new(stdout);

    use rustc_hash::FxHashMap;
    let mut order: Vec<em::ReadProfile> = Vec::new();
    let mut idx: FxHashMap<String, usize> = FxHashMap::default();

    for line in reader.lines() {
        let line = line?;
        let cols: Vec<&str> = line.split('\t').collect();
        if cols.len() < 7 {
            continue;
        }
        let pos: i64 = match cols[1].parse() {
            Ok(p) => p,
            Err(_) => continue,
        };
        // sites=info: only record informative positions (matches em.py `sites=info`)
        let j = match rd.info_idx.get(&pos) {
            Some(&j) => j,
            None => continue,
        };
        let ref_base = cols[2].as_bytes().first().copied().unwrap_or(b'N');
        let alleles = parse_pileup_bases(cols[4].as_bytes(), ref_base);
        let names: Vec<&str> = if cols[6].is_empty() {
            Vec::new()
        } else {
            cols[6].split(',').collect()
        };
        if names.len() != alleles.len() {
            // defensive: token/name mismatch (should not happen); skip column (matches readmodel.py)
            continue;
        }
        for (name, &allele) in names.iter().zip(alleles.iter()) {
            let code = base_code(allele); // ACGT->0..3, N/'*'->-1 (sentinel)
            let entry = *idx.entry((*name).to_string()).or_insert_with(|| {
                order.push(em::ReadProfile::new());
                order.len() - 1
            });
            order[entry].push((j, code));
        }
    }
    let status = child.wait().context("wait samtools mpileup")?;
    anyhow::ensure!(status.success(), "samtools mpileup failed");
    Ok(order)
}

/// Align a raw FASTQ pair to `reference` with strobealign and produce a sorted+indexed BAM. strobealign
/// builds its index on the fly (the compact chrM ref is tiny). This is the non-pure path (external
/// `strobealign` + `samtools`); fast mode stays 100% pure Rust.
fn align_strobealign(
    strobealign: &str,
    samtools: &str,
    reference: &str,
    r1: &str,
    r2: &str,
    out_bam: &str,
    threads: usize,
) -> Result<()> {
    let t = threads.max(1);
    let pipeline = format!(
        "{sb} -t {t} \"{r}\" \"{r1}\" \"{r2}\" | {st} sort -@ {t} -o \"{o}\" -",
        sb = strobealign, st = samtools, r = reference, r1 = r1, r2 = r2, o = out_bam, t = t,
    );
    let status = Command::new("bash")
        .arg("-c")
        .arg(&pipeline)
        .status()
        .context("spawn strobealign | samtools sort pipeline")?;
    anyhow::ensure!(status.success(), "strobealign | samtools sort failed");
    let status = Command::new(samtools)
        .args(["index", out_bam])
        .status()
        .context("spawn samtools index")?;
    anyhow::ensure!(status.success(), "samtools index failed");
    Ok(())
}

/// ACCURATE mode for one sample. Either aligns the raw FASTQ pair with strobealign (when `bam` is
/// None) or uses a caller-supplied pre-aligned BAM, then samtools filter + mpileup -> the same EM.
#[allow(clippy::too_many_arguments)]
fn estimate_accurate(
    rd: &em::RefData,
    reference: &str,
    contig: &str,
    r1: &str,
    r2: &str,
    bam: Option<&str>,
    strobealign: &str,
    samtools: &str,
    sample: &str,
    min_prop: f64,
    n_boot: usize,
    seed: u64,
    threads: usize,
) -> Result<em::RangeEstimate> {
    let safe = sample.replace('/', "_");
    let pid = std::process::id();
    // temp aligned BAM (strobealign path) + temp filtered BAM, both cleaned up at the end.
    let raw = std::env::temp_dir().join(format!("mtnoc_raw_{safe}_{pid}.bam"));
    let raw_s = raw.to_string_lossy().to_string();
    let flt = std::env::temp_dir().join(format!("mtnoc_flt_{safe}_{pid}.bam"));
    let flt_s = flt.to_string_lossy().to_string();

    let result = (|| -> Result<em::RangeEstimate> {
        let in_bam = match bam {
            Some(b) => b.to_string(),
            None => {
                align_strobealign(strobealign, samtools, reference, r1, r2, &raw_s, threads)?;
                raw_s.clone()
            }
        };
        filter_bam_accurate(samtools, &in_bam, &flt_s, contig)?;
        let profiles = pileup_profiles(rd, samtools, reference, contig, &flt_s)?;
        eprintln!("[accurate] fragments: {}", profiles.len());
        // SAME Rust EM as fast mode / bit-exact with Python em.estimate_range.
        Ok(em::estimate_range(rd, &profiles, min_prop, n_boot, 0.8, 8.0, 4000, seed, 200))
    })();

    // best-effort cleanup regardless of outcome
    for p in [&raw_s, &format!("{raw_s}.bai"), &flt_s, &format!("{flt_s}.bai")] {
        let _ = std::fs::remove_file(p);
    }
    result
}


/// A batch of parsed records for one mate, kept as ONE concatenated byte buffer plus per-record
/// end offsets and match flags. This replaces the old per-record `[Vec<u8>;4]` (4 heap allocs per
/// record) with ~2 growing allocations amortised over `CHUNK` records: needletail parses each record
/// zero-copy, we k-mer match on the borrowed `raw_seq()`, and copy the record's full raw bytes
/// (`all()`, the 4 lines minus trailing newline) once into `buf` for the consumer to write in order.
struct Chunk {
    buf: Vec<u8>,     // concatenated record.all() bytes, back to back
    ends: Vec<u32>,   // exclusive end offset in `buf` for each record (start = previous end)
    matched: Vec<bool>,
}
impl Chunk {
    fn with_capacity(n: usize) -> Self {
        Chunk { buf: Vec::with_capacity(n * 384), ends: Vec::with_capacity(n), matched: Vec::with_capacity(n) }
    }
    #[inline]
    fn len(&self) -> usize {
        self.ends.len()
    }
    #[inline]
    fn is_empty(&self) -> bool {
        self.ends.is_empty()
    }
}

const CHUNK: usize = 8192;

/// Producer: decompress `path` (flate2/zlib-rs MultiGzDecoder — pure Rust), parse records with
/// needletail (zero-copy, over the decompressed stream), match each read, send chunks in order.
fn spawn_producer(
    path: String,
    kset: Arc<KmerSet>,
    min_kmers: usize,
) -> (thread::JoinHandle<Result<()>>, std::sync::mpsc::Receiver<Chunk>) {
    let (tx, rx) = sync_channel::<Chunk>(8);
    let h = thread::spawn(move || -> Result<()> {
        let f = File::open(&path).with_context(|| format!("open {path}"))?;
        // Decompress with flate2 (zlib-rs backend). needletail's compression features are OFF, so
        // parse_fastx_reader treats this already-decompressed stream as plain FASTX.
        let dec = MultiGzDecoder::new(BufReader::with_capacity(1 << 20, f));
        let mut reader = parse_fastx_reader(dec).with_context(|| format!("parse {path}"))?;
        let mut chunk = Chunk::with_capacity(CHUNK);
        while let Some(rec) = reader.next() {
            let rec = rec.with_context(|| format!("record in {path}"))?;
            let m = kset.matches(&rec.raw_seq(), min_kmers);
            chunk.buf.extend_from_slice(rec.all());
            chunk.ends.push(chunk.buf.len() as u32);
            chunk.matched.push(m);
            if chunk.len() >= CHUNK {
                if tx.send(std::mem::replace(&mut chunk, Chunk::with_capacity(CHUNK))).is_err() {
                    break;
                }
            }
        }
        if !chunk.is_empty() {
            let _ = tx.send(chunk);
        }
        Ok(())
    });
    (h, rx)
}

/// Parse one raw 4-line FASTQ record (as produced by needletail `rec.all()`, lines joined by `\n`,
/// no trailing newline) into a `myaln::Read`. QNAME = after `@`, up to the first whitespace, commas
/// sanitised to `_` (so it matches the accurate-mode mpileup QNAME handling).
fn parse_fastq_record(raw: &[u8], mate: u8) -> Option<myaln::Read> {
    let mut it = raw.split(|&b| b == b'\n');
    let header = it.next()?;
    let seq = it.next()?;
    let _plus = it.next()?;
    let qual = it.next()?;
    if header.first() != Some(&b'@') {
        return None;
    }
    let nb = &header[1..];
    let end = nb.iter().position(|b| b.is_ascii_whitespace()).unwrap_or(nb.len());
    let name: String = nb[..end].iter().map(|&b| if b == b',' { '_' } else { b as char }).collect();
    Some(myaln::Read { name, mate, seq: seq.to_vec(), qual: qual.to_vec() })
}

/// k-mer prefilter one raw FASTQ pair against `kset`, keeping a pair when EITHER mate shares a
/// canonical 21-mer. Returns the kept reads (both mates, tagged) for the aligner; if `out1`/`out2`
/// are given, the kept raw records are also written there. Two producer threads (one per mate,
/// decompress+match) feed an in-order consumer (gzip is serial-per-stream, so 2 is the useful max).
fn prefilter(
    kset: &Arc<KmerSet>,
    r1: &str,
    r2: &str,
    out1: Option<&str>,
    out2: Option<&str>,
    min_kmers: usize,
) -> Result<Vec<myaln::Read>> {
    let (h1, rx1) = spawn_producer(r1.to_string(), kset.clone(), min_kmers);
    let (h2, rx2) = spawn_producer(r2.to_string(), kset.clone(), min_kmers);

    let mut w1 = out1.map(File::create).transpose()?.map(|f| BufWriter::with_capacity(1 << 20, f));
    let mut w2 = out2.map(File::create).transpose()?.map(|f| BufWriter::with_capacity(1 << 20, f));

    // Keep mate-1 and mate-2 reads in separate vectors and return mate-1s first, then mate-2s. This
    // reproduces the original two-pass order (candidate FASTQ c1 fully, then c2) that the EM's
    // deterministic read-cap subsample depends on — an interleaved order would change the estimate.
    let mut reads1: Vec<myaln::Read> = Vec::new();
    let mut reads2: Vec<myaln::Read> = Vec::new();
    let mut total: u64 = 0;
    let mut kept: u64 = 0;
    loop {
        let c1 = rx1.recv().ok();
        let c2 = rx2.recv().ok();
        match (c1, c2) {
            (Some(c1), Some(c2)) => {
                if c1.len() != c2.len() {
                    anyhow::bail!("R1/R2 chunk length mismatch");
                }
                let (mut s1, mut s2) = (0u32, 0u32);
                for i in 0..c1.len() {
                    let (e1, e2) = (c1.ends[i], c2.ends[i]);
                    total += 1;
                    if c1.matched[i] || c2.matched[i] {
                        kept += 1;
                        let rec1 = &c1.buf[s1 as usize..e1 as usize];
                        let rec2 = &c2.buf[s2 as usize..e2 as usize];
                        if let Some(w) = w1.as_mut() {
                            w.write_all(rec1)?;
                            w.write_all(b"\n")?;
                        }
                        if let Some(w) = w2.as_mut() {
                            w.write_all(rec2)?;
                            w.write_all(b"\n")?;
                        }
                        if let Some(r) = parse_fastq_record(rec1, 1) {
                            reads1.push(r);
                        }
                        if let Some(r) = parse_fastq_record(rec2, 2) {
                            reads2.push(r);
                        }
                    }
                    s1 = e1;
                    s2 = e2;
                }
            }
            (None, None) => break,
            _ => anyhow::bail!("R1/R2 record count mismatch"),
        }
    }
    if let Some(w) = w1.as_mut() {
        w.flush()?;
    }
    if let Some(w) = w2.as_mut() {
        w.flush()?;
    }
    h1.join().unwrap()?;
    h2.join().unwrap()?;
    eprintln!(
        "[fast] prefilter pairs total={total} kept={kept} ({:.4}%)",
        100.0 * kept as f64 / total.max(1) as f64
    );
    reads1.append(&mut reads2);
    Ok(reads1)
}

/// One resolved sample: derived name, paired inputs, and optional per-sample outputs / pre-aligned BAM.
struct Job {
    name: String,
    r1: String,
    r2: String,
    out1: Option<String>,
    out2: Option<String>,
    bam: Option<String>,
}

/// Derive a sample name from an R1 filename: strip the directory and the `_R1`/`.R1`/`_1`/`.1`
/// mate tag (plus everything after it, i.e. the extension).
fn derive_name(r1: &str) -> String {
    let base = std::path::Path::new(r1)
        .file_name()
        .map(|s| s.to_string_lossy().into_owned())
        .unwrap_or_else(|| r1.to_string());
    for tag in ["_R1", ".R1", "_1.", ".1."] {
        if let Some(idx) = base.find(tag) {
            return base[..idx].to_string();
        }
    }
    base.split('.').next().unwrap_or(&base).to_string()
}

/// Resolve `--in`/`--in1`/`--in2` shorthands into concrete R1/R2 paths for one stem. Probes common
/// mate-tag + extension conventions and returns the first pair that exists on disk.
fn resolve_stem(stem: &str) -> Result<(String, String)> {
    for (t1, t2) in [("_R1", "_R2"), (".R1", ".R2"), ("_1", "_2"), (".1", ".2")] {
        for ext in [".fastq.gz", ".fq.gz", ".fastq", ".fq"] {
            let r1 = format!("{stem}{t1}{ext}");
            let r2 = format!("{stem}{t2}{ext}");
            if std::path::Path::new(&r1).exists() && std::path::Path::new(&r2).exists() {
                return Ok((r1, r2));
            }
        }
    }
    anyhow::bail!("could not resolve R1/R2 for stem '{stem}' (tried _R1/.R1/_1/.1 × .fastq[.gz]/.fq[.gz])")
}

/// Build the sample job list from the shared estimate arguments (fast or accurate).
fn resolve_jobs(a: &EstimateArgs) -> Result<Vec<Job>> {
    // 1. primary paired inputs
    let pairs: Vec<(String, String)> = if !a.in_stem.is_empty() {
        a.in_stem.iter().map(|s| resolve_stem(s)).collect::<Result<_>>()?
    } else if !a.in1.is_empty() || !a.in2.is_empty() {
        anyhow::ensure!(
            a.in1.len() == a.in2.len(),
            "--in1 and --in2 must list the same number of files ({} vs {})",
            a.in1.len(),
            a.in2.len()
        );
        a.in1.iter().cloned().zip(a.in2.iter().cloned()).collect()
    } else if !a.bam.is_empty() {
        // accurate with pre-aligned BAM(s) and no FASTQ: name from the BAM, no reads needed
        Vec::new()
    } else {
        anyhow::bail!("no input: give --in1/--in2, --in <stem>, or (accurate) --bam");
    };

    // 2. optional per-sample outputs
    let outs: Vec<(Option<String>, Option<String>)> = if !a.out_stem.is_empty() {
        a.out_stem
            .iter()
            .map(|s| (Some(format!("{s}_R1.fq")), Some(format!("{s}_R2.fq"))))
            .collect()
    } else if !a.out1.is_empty() || !a.out2.is_empty() {
        anyhow::ensure!(
            a.out1.len() == a.out2.len(),
            "--out1 and --out2 must list the same number of files"
        );
        a.out1.iter().cloned().map(Some).zip(a.out2.iter().cloned().map(Some)).collect()
    } else {
        Vec::new()
    };

    // 3. BAM-only accurate path
    if pairs.is_empty() && !a.bam.is_empty() {
        return Ok(a
            .bam
            .iter()
            .map(|b| Job {
                name: std::path::Path::new(b)
                    .file_stem()
                    .map(|s| s.to_string_lossy().into_owned())
                    .unwrap_or_else(|| b.clone()),
                r1: String::new(),
                r2: String::new(),
                out1: None,
                out2: None,
                bam: Some(b.clone()),
            })
            .collect());
    }

    if !outs.is_empty() {
        anyhow::ensure!(
            outs.len() == pairs.len(),
            "number of outputs ({}) must match number of input samples ({})",
            outs.len(),
            pairs.len()
        );
    }
    if !a.bam.is_empty() {
        anyhow::ensure!(
            a.bam.len() == pairs.len(),
            "number of --bam ({}) must match number of input samples ({})",
            a.bam.len(),
            pairs.len()
        );
    }

    Ok(pairs
        .into_iter()
        .enumerate()
        .map(|(i, (r1, r2))| {
            let name = derive_name(&r1);
            let (out1, out2) = outs.get(i).cloned().unwrap_or((None, None));
            Job { name, r1, r2, out1, out2, bam: a.bam.get(i).cloned() }
        })
        .collect())
}

/// Run `fast` or `accurate` over all resolved samples, emitting one TSV row each.
fn run_estimate_cmd(a: &EstimateArgs, accurate: bool) -> Result<()> {
    if a.thread > 0 {
        rayon::ThreadPoolBuilder::new().num_threads(a.thread).build_global().ok();
    }
    let jobs = resolve_jobs(a)?;
    anyhow::ensure!(!jobs.is_empty(), "no samples to process");

    let rd = embedded_refdata()?;
    let index = myaln::RefIndex::build_multi(&a.r#ref, &a.contig)?;
    // fast mode also needs the embedded k-mer prefilter reference
    let kset = if accurate {
        Arc::new(KmerSet::from_seq(&[]))
    } else {
        Arc::new(KmerSet::from_seq(&embedded_kmer_ref()?))
    };
    if !accurate {
        eprintln!("[fast] prefilter reference k-mers: {}", kset.len());
    }

    if a.header {
        println!("{OUT_HEADER}");
    }
    let out = std::io::stdout();
    for job in &jobs {
        eprintln!("[{}] sample {}", if accurate { "accurate" } else { "fast" }, job.name);
        let est = if accurate {
            estimate_accurate(
                &rd, &a.r#ref, &a.contig, &job.r1, &job.r2, job.bam.as_deref(),
                &a.strobealign, &a.samtools, &job.name, a.min_prop, a.n_boot, a.seed, a.thread,
            )?
        } else {
            estimate_fast(
                &rd, &kset, &index, &job.r1, &job.r2, job.out1.as_deref(), job.out2.as_deref(),
                a.min_identity, a.min_prop, a.n_boot, a.seed,
            )?
        };
        use std::io::Write as _;
        writeln!(out.lock(), "{}", format_row(&job.name, &rd, &est))?;
    }
    Ok(())
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.cmd {
        Cmd::Fast(a) => run_estimate_cmd(&a, false)?,
        Cmd::Accurate(a) => run_estimate_cmd(&a, true)?,
        Cmd::Validate { reference, contig, c1, c2, min_identity, truth_names, truth_pos } => {
            use rustc_hash::FxHashMap;
            use std::collections::HashSet;
            let truth: HashSet<String> = BufReader::new(File::open(&truth_names)?)
                .lines()
                .filter_map(|l| l.ok())
                .map(|l| l.trim().to_string())
                .filter(|l| !l.is_empty())
                .collect();
            let mut tpos: FxHashMap<(String, u8), i64> = FxHashMap::default();
            for l in BufReader::new(File::open(&truth_pos)?).lines() {
                let l = l?;
                let mut it = l.split('\t');
                if let (Some(q), Some(m), Some(p)) = (it.next(), it.next(), it.next()) {
                    if let (Ok(m), Ok(p)) = (m.parse::<u8>(), p.parse::<i64>()) {
                        tpos.entry((q.to_string(), m)).and_modify(|e| { if p < *e { *e = p } }).or_insert(p);
                    }
                }
            }
            myaln::run_validate(&reference, &contig, &c1, &c2, min_identity, &truth, &tpos)?
        }
        Cmd::Rngtest => em::rng_selftest(),
        Cmd::Emtest { profiles, min_prop, n_boot, seed } => {
            run_emtest(&profiles, min_prop, n_boot, seed)?
        }
    }
    Ok(())
}