bioai/sequence_utils.py

"""bioai.sequence_utils -- shared sequence helpers (one-hot encode, FASTA IO,
k-mer off-target index, dsRNA tiling).

These are refactored copies of the helpers that previously lived in the
broken `src/` scaffold so that the new `bioai/` package is self-contained.
"""

from __future__ import annotations

from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterator, List, Set, Tuple

import numpy as np

# Canonical safety panel — 12 species covering all ecological roles regulators
# care about: pollinators, beneficial insects, soil invertebrates, aquatic
# invertebrates, livestock (mammals + poultry + fish), and human safety.
# `homo_sapiens` is intentionally last so demo runs without human data still
# produce a column for it (always 0.0).
SAFETY_SPECIES: List[str] = [
    # Pollinators
    "apis_mellifera",        # Western honeybee
    "bombus_terrestris",     # Buff-tailed bumblebee
    "megachile_rotundata",   # Alfalfa leafcutter bee
    # Beneficial predators (biocontrol agents)
    "adalia_bipunctata",     # Two-spot ladybird beetle
    "chrysoperla_carnea",    # Green lacewing
    # Soil invertebrates
    "eisenia_fetida",        # Red wiggler compost worm
    # Aquatic invertebrates
    "daphnia_magna",         # Water flea (standard ecotox model)
    # Livestock & poultry
    "bos_taurus",            # Cattle
    "bos_indicus",           # Zebu
    "gallus_gallus",         # Chicken
    "ovis_aries",            # Sheep
    "sus_scrofa",            # Pig
    # Aquatic vertebrate
    "danio_rerio",           # Zebrafish
    # Human (applicator + consumer safety)
    "homo_sapiens",
]

# Pest species the demo can target. The orchestrator's keyword parser maps
# user input to one of these. Synthetic transcripts are generated for each.
PEST_SPECIES: List[str] = [
    "nilaparvata_lugens",     # Brown planthopper (rice)
    "spodoptera_frugiperda",  # Fall armyworm (maize)
    "schistocerca_gregaria",  # Desert locust (wheat, cereals)
    "chilo_suppressalis",     # Striped stem borer (rice)
    "myzus_persicae",         # Peach-potato aphid (vegetables)
    "leptinotarsa_decemlineata",  # Colorado potato beetle (potato)
    "bemisia_tabaci",         # Tobacco whitefly (tomato, cotton)
]

TARGET_SPECIES: str = "nilaparvata_lugens"

# Common-name -> scientific-name mapping for pest species normalization.
# The LLM (and users) may return "Brown Planthopper" or "brown planthopper"
# but the FASTA headers use "NILAPARVATA_LUGENS_FAKE_001". This mapping
# normalizes any of the accepted aliases to the canonical scientific name
# used in the FASTA headers and the PEST_SPECIES list.
PEST_SPECIES_ALIASES: Dict[str, str] = {
    # Brown planthopper
    "brown planthopper": "nilaparvata_lugens",
    "brown plant hopper": "nilaparvata_lugens",
    "nilaparvata lugens": "nilaparvata_lugens",
    "nilaparvata_lugens": "nilaparvata_lugens",
    "bph": "nilaparvata_lugens",
    # Fall armyworm
    "fall armyworm": "spodoptera_frugiperda",
    "fall army worm": "spodoptera_frugiperda",
    "spodoptera frugiperda": "spodoptera_frugiperda",
    "spodoptera_frugiperda": "spodoptera_frugiperda",
    "armyworm": "spodoptera_frugiperda",
    # Desert locust
    "desert locust": "schistocerca_gregaria",
    "locust": "schistocerca_gregaria",
    "schistocerca gregaria": "schistocerca_gregaria",
    "schistocerca_gregaria": "schistocerca_gregaria",
    # Striped stem borer
    "striped stem borer": "chilo_suppressalis",
    "stem borer": "chilo_suppressalis",
    "chilo suppressalis": "chilo_suppressalis",
    "chilo_suppressalis": "chilo_suppressalis",
    # Peach-potato aphid
    "peach-potato aphid": "myzus_persicae",
    "peach potato aphid": "myzus_persicae",
    "aphid": "myzus_persicae",
    "myzus persicae": "myzus_persicae",
    "myzus_persicae": "myzus_persicae",
    # Colorado potato beetle
    "colorado potato beetle": "leptinotarsa_decemlineata",
    "potato beetle": "leptinotarsa_decemlineata",
    "leptinotarsa decemlineata": "leptinotarsa_decemlineata",
    "leptinotarsa_decemlineata": "leptinotarsa_decemlineata",
    # Tobacco whitefly
    "tobacco whitefly": "bemisia_tabaci",
    "whitefly": "bemisia_tabaci",
    "white fly": "bemisia_tabaci",
    "bemisia tabaci": "bemisia_tabaci",
    "bemisia_tabaci": "bemisia_tabaci",
}


def normalize_pest_species(name: str) -> str:
    """Normalize a pest species name to the canonical scientific name.

    Handles common names ("Brown Planthopper"), scientific names with spaces
    ("Nilaparvata lugens"), and canonical names ("nilaparvata_lugens").
    Returns the input lowercased if no alias matches (defensive).
    """
    if not name:
        return "nilaparvata_lugens"  # default
    key = name.strip().lower()
    if key in PEST_SPECIES_ALIASES:
        return PEST_SPECIES_ALIASES[key]
    # Try removing underscores and spaces
    key_nospace = key.replace("_", " ").replace("-", " ")
    if key_nospace in PEST_SPECIES_ALIASES:
        return PEST_SPECIES_ALIASES[key_nospace]
    # If it's already a canonical name (with underscore), return as-is
    if key in PEST_SPECIES:
        return key
    # Default: return the original lowercased (will likely fail FASTA match,
    # but the caller falls back to using all transcripts)
    return key

# A/C/G/T (U -> T) one-hot mapping used by SiRNACNN.
NUCLEOTIDE_INDEX: Dict[str, int] = {"A": 0, "C": 1, "G": 2, "T": 3, "U": 3}


# --------------------------------------------------------------------------- #
# FASTA I/O
# --------------------------------------------------------------------------- #
def read_fasta(path: str | Path) -> Dict[str, str]:
    """Return ``{sequence_id: sequence}`` from a FASTA file (full load)."""
    path = Path(path)
    seqs: Dict[str, str] = {}
    current_id: str | None = None
    chunks: List[str] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith(">"):
                if current_id is not None:
                    seqs[current_id] = "".join(chunks)
                current_id = line[1:].split()[0]
                chunks = []
            else:
                chunks.append(line.upper())
        if current_id is not None:
            seqs[current_id] = "".join(chunks)
    return seqs


def fasta_iter(path: str | Path) -> Iterator[Tuple[str, str]]:
    """Generator over ``(id, sequence)`` for large FASTA files."""
    path = Path(path)
    with path.open("r", encoding="utf-8") as f:
        current_id: str | None = None
        chunks: List[str] = []
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith(">"):
                if current_id is not None:
                    yield current_id, "".join(chunks)
                current_id = line[1:].split()[0]
                chunks = []
            else:
                chunks.append(line.upper())
        if current_id is not None:
            yield current_id, "".join(chunks)


# --------------------------------------------------------------------------- #
# One-hot encoding
# --------------------------------------------------------------------------- #
def one_hot_encode(seq: str, max_len: int = 21) -> np.ndarray:
    """Channel-first one-hot encoding with shape ``(4, max_len)``.

    Unknown characters get a uniform ``0.25`` distribution over the 4 bases
    so the model still sees a valid probability row instead of all-zeros.
    """
    seq = seq.upper().replace("U", "T")
    encoded = np.zeros((4, max_len), dtype=np.float32)
    for i, ch in enumerate(seq[:max_len]):
        if ch in NUCLEOTIDE_INDEX:
            encoded[NUCLEOTIDE_INDEX[ch], i] = 1.0
        else:
            encoded[:, i] = 0.25
    # pad remaining columns with the uniform distribution too
    for i in range(len(seq), max_len):
        encoded[:, i] = 0.25
    return encoded


def encode_batch(sequences: List[str], max_len: int = 21) -> np.ndarray:
    """Batch of sequences -> numpy array ``(batch, 4, max_len)``."""
    return np.stack([one_hot_encode(s, max_len) for s in sequences], axis=0)


def tokens_to_onehot(tokens: np.ndarray, seq_len: int) -> np.ndarray:
    """Convert ``(batch, seq_len)`` integer tokens in ``{0,1,2,3}`` to
    ``(batch, 4, seq_len)`` one-hot float32.
    """
    batch = tokens.shape[0]
    onehot = np.zeros((batch, 4, seq_len), dtype=np.float32)
    for b in range(batch):
        for i, t in enumerate(tokens[b]):
            if 0 <= int(t) < 4:
                onehot[b, int(t), i] = 1.0
            else:
                onehot[b, :, i] = 0.25
    return onehot


def onehot_to_tokens(onehot: np.ndarray) -> np.ndarray:
    """Inverse of :func:`tokens_to_onehot` -- ``(batch, 4, seq_len)`` ->
    ``(batch, seq_len)`` integer tokens.
    """
    return onehot.argmax(axis=1)


# --------------------------------------------------------------------------- #
# Reverse complement / k-mer generation
# --------------------------------------------------------------------------- #
_COMPLEMENT = {"A": "T", "T": "A", "U": "A", "C": "G", "G": "C"}


def reverse_complement(seq: str) -> str:
    return "".join(_COMPLEMENT.get(b, b) for b in reversed(seq))


def generate_kmers(seq: str, k: int = 21) -> Set[str]:
    seq = seq.upper().replace("U", "T")
    kmers: Set[str] = set()
    if len(seq) < k:
        return kmers
    for i in range(len(seq) - k + 1):
        kmer = seq[i:i + k]
        kmers.add(kmer)
        kmers.add(reverse_complement(kmer))
    return kmers


# --------------------------------------------------------------------------- #
# dsRNA precursor tiling
# --------------------------------------------------------------------------- #
def tile_sequence(
    sequence: str,
    window: int = 200,
    step: int = 100,
    max_candidates: int | None = None,
) -> List[Tuple[int, int, str]]:
    """Tile a transcript into ``window``-nt precursors.

    Returns a list of ``(start, end, subseq)`` tuples. ``step`` defaults to
    ``window // 2`` so precursors overlap by 50% (covers more splice variants
    without blowing up the candidate count).
    """
    seq = sequence.upper().replace("U", "T")
    if step <= 0:
        step = max(1, window // 2)
    candidates: List[Tuple[int, int, str]] = []
    for i in range(0, max(0, len(seq) - window + 1), step):
        if max_candidates is not None and len(candidates) >= max_candidates:
            break
        candidates.append((i, i + window, seq[i:i + window]))
    return candidates


def dice_precursor(precursor: str, sirna_len: int = 21, step: int = 21) -> List[str]:
    """Dice a 200-nt dsRNA precursor into 21-nt siRNAs (Dicer-style).

    Default ``step=21`` matches Dicer's processive 21-nt cut cadence. Use a
    smaller step (e.g. 7) if you want overlapping windows for dense coverage.
    """
    seq = precursor.upper().replace("U", "T")
    sirnas: List[str] = []
    for i in range(0, max(0, len(seq) - sirna_len + 1), step):
        sirnas.append(seq[i:i + sirna_len])
    return sirnas


# --------------------------------------------------------------------------- #
# K-mer off-target index (copied from src/offtarget/kmer_index.py, no edits
# to behaviour -- the API is what the rest of the pipeline expects)
# --------------------------------------------------------------------------- #
class KmerOffTargetIndex:
    """Exact k-mer (default 21-mer, with reverse complements) off-target index.

    ``build_from_fasta`` ingests one species; ``per_species_risk`` returns a
    dict mapping each indexed species -> fraction of the candidate's k-mers
    that hit that species.
    """

    def __init__(self, k: int = 21):
        self.k = k
        self.index: Dict[str, int] = defaultdict(int)
        self.species_kmers: Dict[str, Set[str]] = {}

    def build_from_fasta(
        self,
        fasta_path: str | Path,
        species_name: str,
        header_prefix: str | None = None,
    ) -> None:
        """Ingest one species from a FASTA file.

        If ``header_prefix`` is provided, only sequences whose FASTA header
        starts with ``header_prefix`` are ingested (e.g. ``apis_mellifera``
        matches headers like ``>apis_mellifera_fake_001``). This lets a single
        multi-species FASTA be indexed per-species without pre-splitting.

        If ``header_prefix`` is None (default), all sequences are ingested.
        """
        species_set: Set[str] = set()
        n_seqs = 0
        for header, seq in fasta_iter(fasta_path):
            if header_prefix is not None:
                # Header is e.g. "apis_mellifera_fake_001". Match the prefix.
                if not header.startswith(header_prefix):
                    continue
            n_seqs += 1
            for kmer in generate_kmers(seq, self.k):
                self.index[kmer] += 1
                species_set.add(kmer)
        self.species_kmers[species_name] = species_set
        print(f"Indexed {len(species_set)} unique {self.k}-mers from {species_name} ({n_seqs} seqs)")

    def candidate_offtarget_score(self, candidate_seq: str) -> float:
        cand_kmers = generate_kmers(candidate_seq, self.k)
        if not cand_kmers:
            return 0.0
        hits = sum(1 for kmer in cand_kmers if kmer in self.index)
        return hits / len(cand_kmers)

    def per_species_risk(self, candidate_seq: str) -> Dict[str, float]:
        cand_kmers = generate_kmers(candidate_seq, self.k)
        if not cand_kmers:
            return {sp: 0.0 for sp in self.species_kmers}
        return {
            sp: sum(1 for kmer in cand_kmers if kmer in kmers) / len(cand_kmers)
            for sp, kmers in self.species_kmers.items()
        }