rust-ise 0.2.3

//! ISOSDB ingest — the deep, divergent-lineage side of the union DB.
//!
//! Each ISOSDB entry is a nucleotide IS element (`>ISOSDB<n>`); we translate ALL of
//! its ORFs with the rustygal META gene-finder (in-process, == pyrodigal meta) and
//! tag each ORF with the element's IS family from `IS_fam_annot.txt`. This depth
//! (36k ORFs vs ISfinder's 5k) is what recovers divergent transposases ISfinder's
//! curated references miss (+~8% recall on the metagenome benchmark).
//!
//! In addition to the fam-tagged protein ORFs (the union-profile input), this pass
//! also captures the per-ORF **nucleotide** CDS (rustygal `nuc` output), keyed by
//! the same `ISOSDB<n>_orf<k>` id. Those nt CDS are the IS-lineage POSITIVE pool for
//! the FP-control reference DB (`<out>/fpc/refset`, `P|`-prefixed) — see `build_db::run`.

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

/// Translate every ISOSDB element to fam-tagged transposase ORFs via rustygal meta.
/// ORF ids are normalised to the reference form `ISOSDB<n>_orf<k>`.
///
/// Returns `(orfs, nt_cds)`:
///   * `orfs`   — the fam-tagged **protein** ORFs (union-profile input), MIN_AA-filtered;
///   * `nt_cds` — the matching per-ORF **nucleotide** CDS `(ISOSDB<n>_orf<k>, nt_bytes)`
///                for exactly the same ORF set (IS-lineage FP-control positive pool).
///
/// `nt_cds` is in the same (element, gene-index) order as `orfs`.
pub fn ingest_isosdb(
    fna: &Path,
    fam_annot: &Path,
) -> Result<(Vec<DbOrf>, Vec<(String, Vec<u8>)>), String> {
    let elem_fam = parse_fam_annot(fam_annot)?; // ISOSDB<n> -> family
    let elems = super::read_fasta(fna)?; // (elem_id, nt bytes) in file order
    let meta = rustygal::meta_api::meta_bins();

    let mut out: Vec<DbOrf> = Vec::new();
    let mut nt_cds: Vec<(String, Vec<u8>)> = Vec::new();
    for (i, (elem_id, dna)) in elems.iter().enumerate() {
        let res = rustygal::meta_api::run_meta(i as i32 + 1, elem_id, dna, &meta);
        let family = elem_fam
            .get(elem_id)
            .cloned()
            .unwrap_or_else(|| "unknown".to_string());
        // Protein ORFs (MIN_AA-filtered): these define the kept ORF set.
        let prot_orfs = parse_trans_faa(&res.trans_faa, elem_id);
        // Per-ORF nucleotide CDS, keyed by gene index (same header format, no length filter).
        let nt_by_idx = parse_nuc_faa(&res.nuc, elem_id);
        for (gene_idx, prot) in prot_orfs {
            let id = format!("{elem_id}_orf{gene_idx}");
            if let Some(nt) = nt_by_idx.get(&gene_idx) {
                nt_cds.push((id.clone(), nt.clone()));
            }
            out.push(DbOrf {
                id,
                family: family.clone(),
                source: "ISOSDB",
                seq: prot,
            });
        }
    }
    Ok((out, nt_cds))
}

/// Parse rustygal `trans_faa` bytes (`>{elem}_{geneidx} # b # e # strand # ID=…`)
/// into `(gene_idx, protein_residues)`, stripping a trailing `*` stop if present.
fn parse_trans_faa(data: &[u8], elem_id: &str) -> Vec<(u32, Vec<u8>)> {
    let text = String::from_utf8_lossy(data);
    let mut out: Vec<(u32, Vec<u8>)> = Vec::new();
    let mut idx: Option<u32> = None;
    let mut seq: Vec<u8> = Vec::new();
    let prefix = format!("{elem_id}_");
    for line in text.lines() {
        if let Some(hdr) = line.strip_prefix('>') {
            if let Some(k) = idx.take() {
                push_orf(&mut out, k, std::mem::take(&mut seq));
            }
            seq.clear();
            let tok = hdr.split_whitespace().next().unwrap_or("");
            // tok = "{elem}_{geneidx}" -> parse the geneidx suffix.
            idx = tok
                .strip_prefix(&prefix)
                .and_then(|s| s.parse::<u32>().ok());
        } else if idx.is_some() {
            seq.extend(line.trim().bytes().filter(|&b| b != b'\r' && b != b'*'));
        }
    }
    if let Some(k) = idx.take() {
        push_orf(&mut out, k, seq);
    }
    out
}

/// Parse rustygal `nuc` bytes (same header format as `trans_faa`:
/// `>{elem}_{geneidx} # …`) into `gene_idx -> nucleotide_CDS_bytes`.
/// No length filter here — the caller keeps only the indices that survived the
/// protein MIN_AA filter, so the nt pool tracks the kept ORF set exactly.
fn parse_nuc_faa(data: &[u8], elem_id: &str) -> HashMap<u32, Vec<u8>> {
    let text = String::from_utf8_lossy(data);
    let mut out: HashMap<u32, Vec<u8>> = HashMap::new();
    let mut idx: Option<u32> = None;
    let mut seq: Vec<u8> = Vec::new();
    let prefix = format!("{elem_id}_");
    for line in text.lines() {
        if let Some(hdr) = line.strip_prefix('>') {
            if let Some(k) = idx.take() {
                out.insert(k, std::mem::take(&mut seq));
            }
            seq.clear();
            let tok = hdr.split_whitespace().next().unwrap_or("");
            idx = tok
                .strip_prefix(&prefix)
                .and_then(|s| s.parse::<u32>().ok());
        } else if idx.is_some() {
            seq.extend(line.trim().bytes().filter(|&b| b != b'\r'));
        }
    }
    if let Some(k) = idx.take() {
        out.insert(k, seq);
    }
    out
}

/// Minimum protein length kept (step1_translate.py: `if len(prot) < 30: continue`).
const MIN_AA: usize = 30;

fn push_orf(out: &mut Vec<(u32, Vec<u8>)>, k: u32, seq: Vec<u8>) {
    if seq.len() >= MIN_AA {
        out.push((k, seq));
    }
}

/// Parse `IS_fam_annot.txt` (`IS\tIS_fam` header, then `ISOSDB<n>\t<family>` rows).
fn parse_fam_annot(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 line in data.lines() {
        let mut it = line.split('\t');
        let (Some(elem), Some(fam)) = (it.next(), it.next()) else {
            continue;
        };
        let (elem, fam) = (elem.trim(), fam.trim());
        if elem.is_empty() || elem == "IS" {
            continue; // header / blank
        }
        map.insert(elem.to_string(), fam.to_string());
    }
    Ok(map)
}