bioai/orchestrator.py

"""bioai.orchestrator -- end-to-end biopesticide design pipeline.

Flow:

  1. ``OllamaClient.parse_pest_report(user_text)``         -> pest dict
  2. Load pest transcripts (synthetic by default)
  3. Tile into 200-nt dsRNA precursors (50% overlap)
  4. Dice each precursor into 21-nt siRNAs (Dicer-style)
  5. ``CandidateRanker.rank(sirnas)``                     -> top-N scored
  6. ``OllamaClient.generate_safety_card(...)``            -> markdown card per top-5
  7. ``OllamaClient.generate_regulatory_memo(...)``        -> EPA-style memo

Returns a dict with ``pest_report``, ``candidates`` (top 10), ``safety_cards``,
``regulatory_memo``, ``total_cost_estimate``.

CLI::

    python -m bioai.orchestrator --user-text "Brown planthopper infestation in rice paddy in Tamil Nadu"

If Ollama is not running (or the model isn't pulled), the LLM calls are skipped and the
returned dict contains a ``degraded_mode`` flag and best-effort strings so
the rest of the pipeline (ranking, safety cards from PINN + off-target
index only) still produces useful output.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Dict, List, Optional

from .agent.ollama_client import OllamaClient, DEFAULT_MODEL
from .inference.ranker import CandidateRanker
from .models.sirna_cnn import resolve_device
from .sequence_utils import (
    SAFETY_SPECIES,
    dice_precursor,
    fasta_iter,
    normalize_pest_species,
    read_fasta,
    tile_sequence,
)

# --------------------------------------------------------------------------- #
# Defaults (resolved from bioai.paths so they're portable across machines)
# --------------------------------------------------------------------------- #
from bioai.paths import (  # noqa: E402
    DEFAULT_PEST_FASTA,
    DEFAULT_SAFETY_FASTA,
    SIRNA_CHECKPOINT,
    PINN_CHECKPOINT,
)

# Rough estimate of local compute cost for Llama 3.2 3B on a laptop GPU.
# Local Ollama is effectively free (no API billing), but we still report an
# opportunity-cost estimate based on cloud GPU-equivalent pricing for context.
USD_PER_1K_INPUT_TOKENS = 0.0001
USD_PER_1K_OUTPUT_TOKENS = 0.0001
USD_PER_1K_INPUT_TOKENS = 0.0009
USD_PER_1K_OUTPUT_TOKENS = 0.0009


# --------------------------------------------------------------------------- #
# Orchestrator
# --------------------------------------------------------------------------- #
class BiopesticideOrchestrator:
    """End-to-end pipeline. See module docstring for the flow diagram."""

    def __init__(
        self,
        pest_fasta: Path = DEFAULT_PEST_FASTA,
        safety_fasta: Path = DEFAULT_SAFETY_FASTA,
        sirna_checkpoint: Path = SIRNA_CHECKPOINT,
        pinn_checkpoint: Path = PINN_CHECKPOINT,
        device: str = "auto",
        max_transcripts: int = 5,
        max_precursors_per_transcript: int = 4,
        max_sirnas_per_precursor: int = 9,  # 200 / 21 ~ 9
        ollama_client: Optional[OllamaClient] = None,
    ):
        self.pest_fasta = Path(pest_fasta)
        self.safety_fasta = Path(safety_fasta)
        self.sirna_checkpoint = Path(sirna_checkpoint)
        self.pinn_checkpoint = Path(pinn_checkpoint)
        self.device = resolve_device(device)
        self.max_transcripts = max_transcripts
        self.max_precursors_per_transcript = max_precursors_per_transcript
        self.max_sirnas_per_precursor = max_sirnas_per_precursor

        # Ranker (loads CNN + PINN weights, builds off-target index)
        safety_paths = {sp: self.safety_fasta for sp in SAFETY_SPECIES}
        self.ranker = CandidateRanker(
            safety_fasta_paths=safety_paths,
            sirna_checkpoint=self.sirna_checkpoint,
            pinn_checkpoint=self.pinn_checkpoint,
            device=str(self.device),
        )

        # Ollama client (None in degraded mode if Ollama server is unreachable)
        self.llm = ollama_client
        self.degraded_mode = False
        if self.llm is None:
            try:
                self.llm = OllamaClient()
                print("[orchestrator] Ollama client initialised.")
            except RuntimeError as exc:
                self.degraded_mode = True
                print(f"[orchestrator] DEGRADED MODE -- {exc}. LLM calls will be skipped.")

        # Token accounting (rough)
        self._input_tokens = 0
        self._output_tokens = 0

    # ------------------------------------------------------------------ #
    def _estimate_tokens(self, text: str) -> int:
        # 1 token ~= 4 chars for English text (Llama tokenizer)
        return max(1, len(text) // 4)

    def _record_tokens(self, prompt: str, response: str) -> None:
        self._input_tokens += self._estimate_tokens(prompt)
        self._output_tokens += self._estimate_tokens(response)

    # ------------------------------------------------------------------ #
    def _load_pest_transcripts(self, pest_species: str | None = None) -> Dict[str, str]:
        if not self.pest_fasta.exists():
            print(f"[orchestrator] pest FASTA missing: {self.pest_fasta}")
            return {}
        all_tx = read_fasta(self.pest_fasta)
        # If we know the target pest species, filter transcripts to that species.
        # Synthetic transcript headers look like "NILAPARVATA_LUGENS_FAKE_001".
        if pest_species:
            prefix = pest_species.upper()
            filtered = {k: v for k, v in all_tx.items() if k.startswith(prefix)}
            if filtered:
                items = list(filtered.items())[: self.max_transcripts]
                return dict(items)
            # Fall back to all transcripts if no match (defensive).
            print(f"[orchestrator] no transcripts matched species '{pest_species}'; using all")
        # Take a few transcripts to keep the demo fast.
        items = list(all_tx.items())[: self.max_transcripts]
        return dict(items)

    def _tile_and_dice(self, transcripts: Dict[str, str]) -> List[Dict]:
        """Return a list of ``{precursor, source, sirnas}`` dicts."""
        out: List[Dict] = []
        for gene_id, seq in transcripts.items():
            windows = tile_sequence(
                seq, window=200, step=100,
                max_candidates=self.max_precursors_per_transcript,
            )
            for start, end, precursor in windows:
                sirnas = dice_precursor(precursor, sirna_len=21, step=21)[: self.max_sirnas_per_precursor]
                out.append({
                    "gene_id": gene_id,
                    "precursor_start": start,
                    "precursor_end": end,
                    "precursor_seq": precursor,
                    "sirnas": sirnas,
                })
        return out

    # ------------------------------------------------------------------ #
    def design(self, user_text: str, top_k: int = 10, pest_species_override: str = None) -> Dict:
        """Run the full pipeline. Returns a result dict (see module docstring).

        If ``pest_species_override`` is provided (e.g. "nilaparvata_lugens"),
        the LLM pest-report parsing step is skipped entirely and the species
        is used directly. This is the fast path used by the pest-card UI:
        ~2 seconds end-to-end vs ~15 seconds with LLM parsing.
        """
        self._input_tokens = 0
        self._output_tokens = 0

        # Step 1: parse pest report (or use the override)
        if pest_species_override:
            # Fast path: skip LLM, build pest_report directly from the override
            normalized = normalize_pest_species(pest_species_override)
            # Infer crop from species
            crop_map = {
                "nilaparvata_lugens": "rice",
                "spodoptera_frugiperda": "maize",
                "schistocerca_gregaria": "wheat",
                "chilo_suppressalis": "rice",
                "myzus_persicae": "vegetables",
                "leptinotarsa_decemlineata": "potato",
                "bemisia_tabaci": "tomato",
            }
            pest_report = {
                "pest_species": normalized,
                "crop": crop_map.get(normalized, "unknown"),
                "severity": "moderate",
                "location": "unspecified",
                "notes": f"Direct selection (LLM parsing skipped for speed)",
                "_raw": "",
                "_degraded": False,
            }
            print(f"[orchestrator] pest_species_override provided: '{pest_species_override}' -> '{normalized}' (LLM parsing skipped)")
        elif self.degraded_mode:
            pest_report = self._degraded_pest_report(user_text)
        else:
            try:
                prompt = user_text
                pest_report = self.llm.parse_pest_report(user_text)
                self._record_tokens(prompt, pest_report.get("_raw", ""))
            except Exception as exc:
                print(f"[orchestrator] parse_pest_report failed ({exc!r}); degraded pest report")
                pest_report = self._degraded_pest_report(user_text)

        # Step 2: load transcripts (filter by pest species if known).
        # Normalize the species name first — the LLM may return a common name
        # like "Brown Planthopper" but the FASTA headers use the scientific
        # name "NILAPARVATA_LUGENS_FAKE_001".
        pest_species_raw = pest_report.get("pest_species") if pest_report else None
        pest_species = normalize_pest_species(pest_species_raw) if pest_species_raw else None
        if pest_species and pest_species != pest_species_raw:
            print(f"[orchestrator] normalized pest species: '{pest_species_raw}' -> '{pest_species}'")
            pest_report["pest_species"] = pest_species  # update so downstream uses the normalized name
        transcripts = self._load_pest_transcripts(pest_species=pest_species)
        if not transcripts:
            return {
                "pest_report": pest_report,
                "candidates": [],
                "safety_cards": [],
                "regulatory_memo": "No pest transcripts available; cannot design candidates.",
                "total_cost_estimate": 0.0,
                "degraded_mode": self.degraded_mode,
                "error": "no pest transcripts",
            }

        # Step 3+4: tile + dice
        precursors = self._tile_and_dice(transcripts)
        all_sirnas: List[str] = []
        for p in precursors:
            all_sirnas.extend(p["sirnas"])
        # Deduplicate
        all_sirnas = list(dict.fromkeys(all_sirnas))
        print(f"[orchestrator] {len(transcripts)} transcripts -> "
              f"{len(precursors)} precursors -> {len(all_sirnas)} unique siRNAs")

        # Step 5: rank
        ranked = self.ranker.rank_detailed(all_sirnas, top_k=top_k)
        for r in ranked:
            r["source_gene"] = next(
                (p["gene_id"] for p in precursors if r["sirna_seq"] in p["sirnas"]),
                "unknown",
            )

        # Step 6: safety cards for top 5
        top5 = ranked[:5]
        safety_cards: List[Dict] = []
        for cand in top5:
            if self.degraded_mode:
                card = self._degraded_safety_card(cand)
            else:
                try:
                    card = self.llm.generate_safety_card(
                        sirna_seq=cand["sirna_seq"],
                        offtarget_risks=cand["offtarget_per_species"],
                        half_life_hours=cand["half_life_hours"],
                    )
                    self._record_tokens(cand["sirna_seq"], card)
                except Exception as exc:
                    print(f"[orchestrator] generate_safety_card failed ({exc!r}); degraded card")
                    card = self._degraded_safety_card(cand)
            safety_cards.append({
                "sirna_seq": cand["sirna_seq"],
                "card_markdown": card,
            })

        # Step 7: regulatory memo
        if self.degraded_mode:
            memo = self._degraded_regulatory_memo(pest_report, ranked[:5])
        else:
            try:
                pest_name = (
                    pest_report.get("pest_species")
                    or "nilaparvata_lugens (brown planthopper)"
                )
                memo = self.llm.generate_regulatory_memo(pest_name, ranked[:5])
                # rough token accounting: prompt ~ all candidate strings
                prompt_blob = pest_name + "".join(
                    c.get("sirna_seq", "") for c in ranked[:5]
                )
                self._record_tokens(prompt_blob, memo)
            except Exception as exc:
                print(f"[orchestrator] generate_regulatory_memo failed ({exc!r}); degraded memo")
                memo = self._degraded_regulatory_memo(pest_report, ranked[:5])

        cost = (
            self._input_tokens * USD_PER_1K_INPUT_TOKENS
            + self._output_tokens * USD_PER_1K_OUTPUT_TOKENS
        ) / 1000.0

        return {
            "pest_report": pest_report,
            "candidates": ranked,
            "safety_cards": safety_cards,
            "regulatory_memo": memo,
            "total_cost_estimate": cost,
            "degraded_mode": self.degraded_mode,
            "n_transcripts": len(transcripts),
            "n_precursors": len(precursors),
            "n_sirnas": len(all_sirnas),
        }

    # ------------------------------------------------------------------ #
    # Degraded-mode helpers (used when Ollama is not running)
    # ------------------------------------------------------------------ #
    @staticmethod
    def _degraded_pest_report(user_text: str) -> Dict:
        """Heuristic pest-species guesser for offline mode.

        We don't try to be smart -- just match a few common pest/crop keywords
        so the demo output looks plausible. The real pipeline uses the LLM.
        """
        text = user_text.lower()
        pest = "nilaparvata_lugens"
        crop = "rice"
        if "aphid" in text:
            pest, crop = "myzus_persicae", "vegetables"
        if "fall armyworm" in text or "spodoptera" in text:
            pest, crop = "spodoptera_frugiperda", "maize"
        if "planthopper" in text or "lugens" in text:
            pest, crop = "nilaparvata_lugens", "rice"
        if "stem borer" in text or "chilo" in text:
            pest, crop = "chilo_suppressalis", "rice"
        if "locust" in text or "schistocerca" in text:
            pest, crop = "schistocerca_gregaria", "wheat"
        if "colorado potato" in text or "leptinotarsa" in text:
            pest, crop = "leptinotarsa_decemlineata", "potato"
        if "whitefly" in text or "bemisia" in text:
            pest, crop = "bemisia_tabaci", "tomato"
        return {
            "pest_species": pest,
            "crop": crop,
            "severity": "moderate",
            "location": "Tamil Nadu, India" if "tamil" in text else "unspecified",
            "notes": "DEGRADED MODE: parsed without LLM (Ollama server not running).",
            "_raw": "",
            "_degraded": True,
        }

    @staticmethod
    def _degraded_safety_card(cand: Dict) -> str:
        """Plain-markdown safety card built from the ranker outputs only."""
        ot_lines = "\n".join(
            f"  - {sp}: {risk:.3f}"
            for sp, risk in cand.get("offtarget_per_species", {}).items()
        ) or "  - (no off-target hits)"
        risk_tier = (
            "high" if cand.get("offtarget_max", 0) > 0.3 else
            "moderate" if cand.get("offtarget_max", 0) > 0.05 else
            "low"
        )
        return (
            f"# Safety Card (DEGRADED MODE)\n\n"
            f"## Sequence\n"
            f"`{cand['sirna_seq']}`\n\n"
            f"## Off-Target Profile\n"
            f"{ot_lines}\n\n"
            f"Max off-target risk: **{cand.get('offtarget_max', 0):.3f}**\n\n"
            f"## Environmental Fate\n"
            f"Predicted half-life: **{cand.get('half_life_hours', 0):.2f} hours**\n\n"
            f"## Overall Risk Tier\n"
            f"**{risk_tier}**\n\n"
            f"_Generated without LLM (Ollama server not running)._"
        )

    @staticmethod
    def _degraded_regulatory_memo(pest_report: Dict, candidates: List[Dict]) -> str:
        pest = pest_report.get("pest_species", "unknown_pest")
        crop = pest_report.get("crop", "unknown_crop")
        lines = [
            f"# Regulatory Memo (DEGRADED MODE)\n",
            f"## Pest & Crop\n",
            f"Target: **{pest}** on **{crop}**.\n",
            f"## Candidate Summary\n",
        ]
        for i, c in enumerate(candidates, 1):
            lines.append(
                f"{i}. `{c['sirna_seq']}`  "
                f"efficacy={c.get('efficacy', 0):.3f}  "
                f"offtarget_max={c.get('offtarget_max', 0):.3f}  "
                f"half_life={c.get('half_life_hours', 0):.2f}h  "
                f"score={c.get('final_score', 0):.3f}"
            )
        lines.append("\n## Risk Assessment\n")
        any_high = any(c.get("offtarget_max", 0) > 0.3 for c in candidates)
        any_short = any(c.get("half_life_hours", 99) < 6 for c in candidates)
        if any_high:
            lines.append("- At least one candidate has HIGH off-target risk; flag for further screening.")
        if any_short:
            lines.append("- At least one candidate has a predicted half-life under 6 hours; field efficacy may be limited.")
        lines.append("\n## Recommendation\n")
        if any_high:
            lines.append("Recommend additional off-target screening before issuing an Experimental Use Permit.")
        else:
            lines.append("Candidates look suitable for an Experimental Use Permit application, pending wet-lab validation.")
        lines.append("\n_Generated without LLM (Ollama server not running)._")
        return "\n".join(lines)


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(description="Run the end-to-end biopesticide design pipeline.")
    p.add_argument("--user-text", type=str, required=True,
                   help="Free-text pest report (e.g. 'Brown planthopper infestation in rice paddy in Tamil Nadu').")
    p.add_argument("--device", type=str, default="auto", choices=["auto", "cpu", "cuda"])
    p.add_argument("--pest-fasta", type=str, default=str(DEFAULT_PEST_FASTA))
    p.add_argument("--safety-fasta", type=str, default=str(DEFAULT_SAFETY_FASTA))
    p.add_argument("--sirna-checkpoint", type=str, default=str(SIRNA_CHECKPOINT))
    p.add_argument("--pinn-checkpoint", type=str, default=str(PINN_CHECKPOINT))
    p.add_argument("--top-k", type=int, default=10)
    p.add_argument("--max-transcripts", type=int, default=5)
    p.add_argument("--pest-species", type=str, default=None,
                   help="skip LLM parsing and use this species directly (e.g. nilaparvata_lugens)")
    args = p.parse_args(argv)

    orch = BiopesticideOrchestrator(
        pest_fasta=Path(args.pest_fasta),
        safety_fasta=Path(args.safety_fasta),
        sirna_checkpoint=Path(args.sirna_checkpoint),
        pinn_checkpoint=Path(args.pinn_checkpoint),
        device=args.device,
        max_transcripts=args.max_transcripts,
    )
    result = orch.design(args.user_text, top_k=args.top_k, pest_species_override=args.pest_species)

    # Pretty-print to stdout
    print("\n" + "=" * 78)
    print("DESIGN RESULT")
    print("=" * 78)
    print(f"Pest species : {result['pest_report'].get('pest_species')}")
    print(f"Crop         : {result['pest_report'].get('crop')}")
    print(f"Severity     : {result['pest_report'].get('severity')}")
    print(f"Location     : {result['pest_report'].get('location')}")
    print(f"Transcripts  : {result['n_transcripts']}")
    print(f"Precursors   : {result['n_precursors']}")
    print(f"siRNAs       : {result['n_sirnas']}")
    print(f"Degraded mode: {result['degraded_mode']}")
    print(f"Cost estimate: ${result['total_cost_estimate']:.4f}")
    print()
    print(f"Top {len(result['candidates'])} candidates:")
    for i, c in enumerate(result["candidates"], 1):
        print(
            f"  {i:2d}. {c['sirna_seq']}  eff={c['efficacy']:.3f}  "
            f"ot_max={c['offtarget_max']:.3f}  hl={c['half_life_hours']:.2f}h  "
            f"score={c['final_score']:.3f}"
        )
    print()
    print("Safety cards (top 5):")
    for sc in result["safety_cards"]:
        print(f"--- {sc['sirna_seq']} ---")
        print(sc["card_markdown"])
        print()
    print("Regulatory memo:")
    print(result["regulatory_memo"])
    return 0


if __name__ == "__main__":
    sys.exit(main())