rust-ise 0.2.3

//! ISfinder ingest — the clean, authoritative side of the union DB.
//!
//! `IS.faa` headers look like:
//!   >html.2020//<ISname> ~~~<orf_descriptor>~~~<function>~~~
//! We keep only ORFs whose `<function>` is exactly `Transposase` (7,121 of them;
//! the `Accessory`/`Passenger` genes are dropped — this is why the ISfinder side
//! carries no passenger-gene contamination). Family comes from `IS.csv` (Name→Family).

use super::DbOrf;
use std::collections::HashMap;
use std::path::Path;

/// Ingest ISfinder transposase ORFs, family-tagged from `IS.csv`.
pub fn ingest_isfinder(is_faa: &Path, is_csv: &Path) -> Result<Vec<DbOrf>, String> {
    let fam = parse_is_csv(is_csv)?; // ISname -> family
    let data = std::fs::read(is_faa).map_err(|e| format!("read {}: {e}", is_faa.display()))?;
    let text = String::from_utf8_lossy(&data);

    let mut out: Vec<DbOrf> = Vec::new();
    let mut keep = false;
    let mut id = String::new();
    let mut family = String::new();
    let mut seq: Vec<u8> = Vec::new();

    for line in text.lines() {
        if let Some(hdr) = line.strip_prefix('>') {
            // Flush the previous record before starting a new one.
            if keep && !seq.is_empty() {
                out.push(DbOrf {
                    id: std::mem::take(&mut id),
                    family: std::mem::take(&mut family),
                    source: "ISfinder",
                    seq: std::mem::take(&mut seq),
                });
            }
            seq.clear();

            let parts: Vec<&str> = hdr.split("~~~").collect();
            let first_tok = parts
                .first()
                .and_then(|s| s.split_whitespace().next())
                .unwrap_or("");
            let func = parts.get(2).map(|s| s.trim()).unwrap_or("");
            keep = func == "Transposase";
            // Unique id = the ORF descriptor (parts[1]); multiple transposase ORFs can
            // share one ISname, so the descriptor keeps profile-sequence names unique.
            let descr = parts
                .get(1)
                .map(|s| s.trim())
                .filter(|s| !s.is_empty())
                .unwrap_or(first_tok);
            id = descr.to_string();
            let isname = first_tok.strip_prefix("html.2020//").unwrap_or(first_tok);
            family = fam.get(isname).cloned().unwrap_or_else(|| "unknown".to_string());
        } else if keep {
            seq.extend(line.trim().bytes().filter(|&b| b != b'\r'));
        }
    }
    if keep && !seq.is_empty() {
        out.push(DbOrf { id, family, source: "ISfinder", seq });
    }
    Ok(out)
}

/// Parse `IS.csv` (header row then rows) into `Name → Family`. Columns 2 (Name) and
/// 3 (Family) precede any quoted field, so a plain comma split is safe for them.
fn parse_is_csv(path: &Path) -> Result<HashMap<String, String>, String> {
    let data = std::fs::read_to_string(path).map_err(|e| format!("read {}: {e}", path.display()))?;
    let mut map = HashMap::new();
    for (i, line) in data.lines().enumerate() {
        if i == 0 || line.trim().is_empty() {
            continue; // header
        }
        let cols: Vec<&str> = line.splitn(4, ',').collect();
        if cols.len() >= 3 {
            let name = cols[1].trim();
            let family = cols[2].trim();
            if !name.is_empty() && !family.is_empty() {
                map.insert(name.to_string(), family.to_string());
            }
        }
    }
    Ok(map)
}

// ================================================================ ISfinder nt CDS
// Auto-derive the ISfinder nucleotide transposase CDS (the ISfinder half of the
// FP-control `P|` positive pool), in pure Rust, from the ISfinder protein records
// (`IS.faa`) + element DNA (`IS.fna`). ISfinder ships no CDS coordinates, so we
// locate each transposase protein inside its element DNA by 6-frame standard-table
// translation, then map the aa hit back to the full coding nt span.
//
// Faithful port of `hth_groups/full/extract_hth_nt_isfinder.py` (primary, exact
// single-frame match) + `recover_isfinder.py` (fallback, longest-core-prefix frame),
// EXCEPT it emits the full transposase CDS `coding = s[nt0 .. nt0+len(prot)*3]` — the
// shipped `fpc/refset` P| pool is full-CDS-scale (p50 ~1020 nt), not the HTH sub-span.
// Determinism: ISfinder records are processed in `IS.faa` file order (first-wins per
// `html.2020//<ISname>` id, mirroring the reference `union_orfs.faa` `setdefault`).

/// Standard NCBI codon table residues, indexed by `(b1*16 + b2*4 + b3)` over base
/// order `TCAG` (internal residues are identical to translation table 11).
const CODON_AA: &[u8; 64] = b"FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG";

/// Base → index in `TCAG` (the codon-table ordering). Non-ACGT → `None` (→ 'X').
fn base_idx(b: u8) -> Option<usize> {
    match b {
        b'T' | b't' => Some(0),
        b'C' | b'c' => Some(1),
        b'A' | b'a' => Some(2),
        b'G' | b'g' => Some(3),
        _ => None,
    }
}

/// Translate one nt frame to residues (`for i in 0..len-2 step 3`), unknown codon → 'X'.
fn tr(s: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(s.len() / 3);
    let mut i = 0;
    while i + 3 <= s.len() {
        let aa = match (base_idx(s[i]), base_idx(s[i + 1]), base_idx(s[i + 2])) {
            (Some(a), Some(b), Some(c)) => CODON_AA[a * 16 + b * 4 + c],
            _ => b'X',
        };
        out.push(aa);
        i += 3;
    }
    out
}

/// Reverse-complement (ACGT/acgt, N passthrough; anything else unchanged, as Python).
fn rc(s: &[u8]) -> Vec<u8> {
    s.iter()
        .rev()
        .map(|&b| match b {
            b'A' => b'T', b'C' => b'G', b'G' => b'C', b'T' => b'A',
            b'a' => b't', b'c' => b'g', b'g' => b'c', b't' => b'a',
            x => x,
        })
        .collect()
}

/// Index of the first occurrence of `needle` in `hay` (like Python `str.find`).
fn find_sub(hay: &[u8], needle: &[u8]) -> Option<usize> {
    if needle.is_empty() {
        return Some(0);
    }
    if needle.len() > hay.len() {
        return None;
    }
    hay.windows(needle.len()).position(|w| w == needle)
}

/// Emit `s[nt0 .. nt0+len]`, truncating at the sequence end (Python slice semantics).
fn slice_coding(s: &[u8], nt0: usize, len: usize) -> Vec<u8> {
    let end = (nt0 + len).min(s.len());
    if nt0 >= s.len() {
        Vec::new()
    } else {
        s[nt0..end].to_vec()
    }
}

/// PRIMARY pass: exact single-frame match of `core` (= prot[1:], `*`-stripped) across
/// both strands / 3 frames; on the first hit, return the full coding nt
/// `s[nt0 .. nt0+len(prot)*3]` where `nt0 = fr + (pos-1)*3`. `None` if not located.
fn primary_locate(seq: &[u8], prot_len: usize, core: &[u8]) -> Option<Vec<u8>> {
    let rev = rc(seq);
    for s in [seq, rev.as_slice()] {
        for fr in 0..3 {
            let aa = tr(&s[fr..]);
            if let Some(pos) = find_sub(&aa, core) {
                // aa index of ORF residue 1 = pos-1 (core starts at prot[1]).
                let nt0 = fr as isize + (pos as isize - 1) * 3;
                if nt0 < 0 {
                    continue; // upstream of the element start — cannot map (rare edge case).
                }
                return Some(slice_coding(s, nt0 as usize, prot_len * 3));
            }
        }
    }
    None
}

/// FALLBACK (recover): for ORFs the primary pass misses, take the strand/frame whose
/// aligned prefix of `core` is longest (binary-search the largest `m` with `core[:m]`
/// present), require `m >= 10`, and emit the full coding anchored at that frame. Ports
/// `recover_isfinder.py`'s longest-prefix selection; adapted to emit the full CDS
/// (not the HTH sub-span), so its translation may diverge past a frameshift — the nt
/// is still real element DNA anchored at the located start.
fn recover_locate(seq: &[u8], prot_len: usize, core: &[u8]) -> Option<Vec<u8>> {
    let rev = rc(seq);
    let mut best_k = 0usize;
    let mut best: Option<(Vec<u8>, usize, usize)> = None; // (frame_seq, fr, pos)
    for s in [seq, rev.as_slice()] {
        for fr in 0..3 {
            let aa = tr(&s[fr..]);
            let (mut lo, mut hi) = (0usize, core.len());
            while lo < hi {
                let m = (lo + hi + 1) / 2;
                if find_sub(&aa, &core[..m]).is_some() {
                    lo = m;
                } else {
                    hi = m - 1;
                }
            }
            if lo > best_k {
                best_k = lo;
                let pos = find_sub(&aa, &core[..lo]).unwrap_or(0);
                best = Some((s.to_vec(), fr, pos));
            }
        }
    }
    if best_k < 10 {
        return None;
    }
    let (s, fr, pos) = best?;
    let nt0 = fr as isize + (pos as isize - 1) * 3;
    if nt0 < 0 {
        return None;
    }
    Some(slice_coding(&s, nt0 as usize, prot_len * 3))
}

/// Counts from an [`isfinder_nt_pool`] run (mirrors the Python tallies).
#[derive(Default, Debug)]
pub struct IsfNtStats {
    /// Unique `html.2020//<ISname>` candidates (first-wins per id, transposase only).
    pub candidates: usize,
    /// Located by the exact single-frame PRIMARY pass (`src=isfinder`).
    pub primary: usize,
    /// Located only by the longest-prefix RECOVER fallback (`src=isfinder_recovered`).
    pub recovered: usize,
    /// Skipped: no matching `IS.fna` element DNA for the ORF's ISname.
    pub no_nt: usize,
    /// Skipped: protein not locatable in any frame (primary and fallback both failed).
    pub not_located: usize,
}

/// Derive the ISfinder nt CDS pool for the FP-control `P|` positive pool from
/// `IS.faa` (transposase proteins) + `IS.fna` (element DNA). Returns
/// `(records, stats)` where each record is `(orf_id = "html.2020//<ISname>", coding_nt)`
/// in `IS.faa` file order (first-wins per id). Deterministic and pure Rust.
pub fn isfinder_nt_pool(
    is_faa: &Path,
    is_fna: &Path,
) -> Result<(Vec<(String, Vec<u8>)>, IsfNtStats), String> {
    // 1. transposase protein per html-id (first_tok), first-wins — mirrors union_orfs.faa.
    let prots = transposase_prot_first_wins(is_faa)?;
    // 2. IS.fna: ISname (first token before '_') -> element DNA (upper), first-wins.
    let fna = load_isfna(is_fna)?;

    let mut out: Vec<(String, Vec<u8>)> = Vec::with_capacity(prots.len());
    let mut st = IsfNtStats { candidates: prots.len(), ..Default::default() };
    for (id, prot) in &prots {
        // name = part after the last "//" (the ISname); IS.fna is keyed by ISname.
        let name = id.rsplit("//").next().unwrap_or(id);
        let seq = match fna.get(name) {
            Some(s) => s,
            None => {
                st.no_nt += 1;
                continue;
            }
        };
        // core = prot[1:] with '*' removed (drop leading M for alt-start codons).
        let core: Vec<u8> = prot.iter().skip(1).filter(|&&b| b != b'*').cloned().collect();
        if core.is_empty() {
            st.not_located += 1;
            continue;
        }
        if let Some(coding) = primary_locate(seq, prot.len(), &core) {
            if !coding.is_empty() {
                out.push((id.clone(), coding));
                st.primary += 1;
                continue;
            }
        }
        if let Some(coding) = recover_locate(seq, prot.len(), &core) {
            if !coding.is_empty() {
                out.push((id.clone(), coding));
                st.recovered += 1;
                continue;
            }
        }
        st.not_located += 1;
    }
    Ok((out, st))
}

/// Read `IS.faa`, returning `(html.2020//<ISname>, protein)` for each **transposase**
/// record, first-wins per id, in file order (reproduces the ISfinder half of the
/// reference `union_orfs.faa`). Protein bytes are uppercased.
fn transposase_prot_first_wins(is_faa: &Path) -> Result<Vec<(String, Vec<u8>)>, String> {
    let data = std::fs::read(is_faa).map_err(|e| format!("read {}: {e}", is_faa.display()))?;
    let text = String::from_utf8_lossy(&data);
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    let mut out: Vec<(String, Vec<u8>)> = Vec::new();
    let mut keep = false;
    let mut id = String::new();
    let mut seq: Vec<u8> = Vec::new();

    for line in text.lines() {
        if let Some(hdr) = line.strip_prefix('>') {
            if keep && !seq.is_empty() && !seen.contains(id.as_str()) {
                seen.insert(id.clone());
                out.push((std::mem::take(&mut id), std::mem::take(&mut seq)));
            }
            seq.clear();
            let parts: Vec<&str> = hdr.split("~~~").collect();
            let first_tok = parts
                .first()
                .and_then(|s| s.split_whitespace().next())
                .unwrap_or("");
            let func = parts.get(2).map(|s| s.trim()).unwrap_or("");
            keep = func == "Transposase";
            id = first_tok.to_string();
        } else if keep {
            seq.extend(line.trim().bytes().filter(|&b| b != b'\r').map(|b| b.to_ascii_uppercase()));
        }
    }
    if keep && !seq.is_empty() && !seen.contains(id.as_str()) {
        out.push((id, seq));
    }
    Ok(out)
}

/// Parse `IS.fna` into `ISname -> element_DNA` (uppercase), first-wins. The key is the
/// header's first whitespace token truncated at the first `_` (`>IS1096_ISL3_unknown`
/// → `IS1096`), matching the reference loader.
fn load_isfna(path: &Path) -> Result<HashMap<String, Vec<u8>>, String> {
    let data = std::fs::read(path).map_err(|e| format!("read {}: {e}", path.display()))?;
    let mut map: HashMap<String, Vec<u8>> = HashMap::new();
    let mut cur: Option<String> = None;
    let mut buf: Vec<u8> = Vec::new();
    for line in data.split(|&b| b == b'\n') {
        if line.first() == Some(&b'>') {
            if let Some(k) = cur.take() {
                map.entry(k).or_insert_with(|| std::mem::take(&mut buf));
                buf.clear();
            }
            let hdr = String::from_utf8_lossy(&line[1..]);
            let tok = hdr.split_whitespace().next().unwrap_or("");
            let key = tok.split('_').next().unwrap_or(tok).to_string();
            cur = Some(key);
        } else {
            buf.extend(line.iter().filter(|&&b| b != b'\r').map(|b| b.to_ascii_uppercase()));
        }
    }
    if let Some(k) = cur.take() {
        map.entry(k).or_insert(buf);
    }
    Ok(map)
}

#[cfg(test)]
mod nt_tests {
    use super::*;
    use std::path::PathBuf;

    // Unit checks for the pure-Rust codon/translation helpers (no external files).
    #[test]
    fn translate_and_rc_basics() {
        assert_eq!(tr(b"ATGAAA"), b"MK"); // ATG=M, AAA=K
        assert_eq!(tr(b"TTTTGA"), b"F*"); // TTT=F, TGA=stop
        assert_eq!(tr(b"ATGA"), b"M"); // trailing partial codon dropped
        assert_eq!(rc(b"ATGC"), b"GCAT");
        assert_eq!(find_sub(b"XXMKQAR", b"MKQ"), Some(2));
        assert_eq!(find_sub(b"XXMKQAR", b"ZZZ"), None);
    }

    // Round-trip: build a synthetic element whose forward frame 0 contains a known ORF,
    // and confirm primary_locate recovers the full coding starting at the (alt) start codon.
    #[test]
    fn primary_locate_recovers_full_cds() {
        // prot = M K Q A R  (len 5). core = "KQAR". CDS codons: ATG AAA CAA GCA CGT.
        let cds = b"ATGAAACAAGCACGT";
        // embed with 6nt upstream + 9nt downstream noise on the forward strand
        let elem = b"GGGCCCATGAAACAAGCACGTTTTGGGAAA";
        let prot = b"MKQAR";
        let core: Vec<u8> = prot.iter().skip(1).cloned().collect();
        let got = primary_locate(elem, prot.len(), &core).expect("located");
        assert_eq!(&got, cds, "full coding == the embedded CDS, start codon first");
    }

    // Full byte-parity check vs the shipped fpc refset's ISfinder P| entries.
    // Gated on the real ISfinder sources + a shipped refset + mmseqs on PATH (CI-safe skip).
    #[test]
    fn isfinder_nt_matches_shipped_refset() {
        // Dev-only integration check — gated on env-supplied paths so no developer
        // absolute path is baked into the published source:
        //   RUST_ISE_TEST_ISFINDER_DIR   → dir holding IS.faa + IS.fna
        //   RUST_ISE_TEST_SHIPPED_REFSET → shipped mmseqs fpc refset db prefix (…/fpc/refset)
        // plus `mmseqs` on PATH. Any unset/missing → skip (CI-safe).
        let dir = match std::env::var_os("RUST_ISE_TEST_ISFINDER_DIR") {
            Some(d) => PathBuf::from(d),
            None => { eprintln!("skip: RUST_ISE_TEST_ISFINDER_DIR unset"); return; }
        };
        let refset = match std::env::var_os("RUST_ISE_TEST_SHIPPED_REFSET") {
            Some(r) => PathBuf::from(r),
            None => { eprintln!("skip: RUST_ISE_TEST_SHIPPED_REFSET unset"); return; }
        };
        let is_faa = dir.join("IS.faa");
        let is_fna = dir.join("IS.fna");
        if !is_faa.exists() || !is_fna.exists() || !refset.with_extension("dbtype").exists() {
            eprintln!("skip: ISfinder sources or shipped refset not present");
            return;
        }
        // Derive the ISfinder nt pool in pure Rust.
        let (recs, st) = isfinder_nt_pool(&is_faa, &is_fna).expect("derive isfinder nt");
        eprintln!(
            "derived ISfinder nt: {} ({} primary + {} recovered) of {} candidates; no_nt={} not_located={}",
            recs.len(), st.primary, st.recovered, st.candidates, st.no_nt, st.not_located
        );

        // Dump the shipped refset to FASTA and collect the P|html entries.
        let tmp = std::env::temp_dir().join(format!("isf_shipped_{}.fasta", std::process::id()));
        let ok = std::process::Command::new("mmseqs")
            .args(["convert2fasta",
                   refset.to_str().unwrap(),
                   tmp.to_str().unwrap(), "-v", "0"])
            .status().map(|s| s.success()).unwrap_or(false);
        if !ok {
            eprintln!("skip: mmseqs convert2fasta unavailable");
            return;
        }
        let data = std::fs::read(&tmp).unwrap();
        let _ = std::fs::remove_file(&tmp);
        let mut shipped_map: HashMap<String, Vec<u8>> = HashMap::new();
        let text = String::from_utf8_lossy(&data);
        let (mut cur, mut buf): (Option<String>, Vec<u8>) = (None, Vec::new());
        for line in text.lines() {
            if let Some(h) = line.strip_prefix('>') {
                if let Some(k) = cur.take() { shipped_map.insert(k, std::mem::take(&mut buf)); }
                let tok = h.split_whitespace().next().unwrap_or("");
                // keep only P|html.* entries, store under the bare orf id (strip "P|")
                cur = tok.strip_prefix("P|").filter(|s| s.starts_with("html")).map(|s| s.to_string());
            } else if cur.is_some() {
                buf.extend(line.trim().bytes());
            }
        }
        if let Some(k) = cur.take() { shipped_map.insert(k, buf); }
        eprintln!("shipped P|html entries: {}", shipped_map.len());

        // Compare: derived vs shipped, by id.
        let (mut exact, mut present, mut diff, mut extra) = (0usize, 0usize, 0usize, 0usize);
        for (id, seq) in &recs {
            match shipped_map.get(id) {
                Some(s) if s == seq => { exact += 1; present += 1; }
                Some(_) => { diff += 1; present += 1; }
                None => extra += 1,
            }
        }
        let missing = shipped_map.len().saturating_sub(present);
        eprintln!(
            "byte-parity vs shipped P|html: exact={exact} diff={diff} \
             derived_extra(not in shipped)={extra} shipped_missing(in shipped not derived)={missing}"
        );
        // The shipped pool is primary-pass-only; expect the primary set to reproduce it
        // to a very high byte-identical rate.
        let rate = exact as f64 / shipped_map.len().max(1) as f64;
        eprintln!("exact byte-match rate vs shipped: {:.1}%", rate * 100.0);
        assert!(rate > 0.95, "≥95% of shipped P|html entries reproduced byte-identically (got {:.1}%)", rate * 100.0);
    }
}