bioai/ui/gradio_app.py

"""
Biopesticide-AI Gradio UI v2 -- high-end startup-grade interface.

Redesigned with:
- Hero section with project pitch and key metrics
- Tabbed workflow: Design | Analytics | Safety | Regulatory | About
- Data visualizations (efficacy chart, off-target heatmap, half-life chart)
- Candidate cards instead of plain tables
- CSV/JSON export buttons
- Professional color scheme and typography
- Loading states and empty states
- Backend status strip with live model/LLM info

Usage:
    python -m bioai.ui.gradio_app            # launch on 0.0.0.0:7860
    python -m bioai.ui.gradio_app --port 8080
    python -m bioai.ui.gradio_app --share     # public share link
"""

from __future__ import annotations

import argparse
import csv
import io
import json
import os
import sys
import time
from pathlib import Path

import gradio as gr

# Make sure we can import bioai from anywhere
_PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from bioai.orchestrator import BiopesticideOrchestrator  # noqa: E402
from bioai.sequence_utils import SAFETY_SPECIES, PEST_SPECIES  # noqa: E402
from bioai.ui.charts import efficacy_bar_chart, offtarget_heatmap, halflife_chart  # noqa: E402

# ─────────────────────────────────────────────────────────────────────────────
# Branding
# ─────────────────────────────────────────────────────────────────────────────
TITLE = "Biopesticide-AI"
TAGLINE = "Design species-specific dsRNA biopesticides in minutes, not months"
SUBTITLE = "Local Llama 3.2 3B + PyTorch + 14-species safety panel + physics-informed fate model"

EXAMPLES = [
    ["Brown planthopper infestation in my rice paddy near Coimbatore, Tamil Nadu. Severity moderate, second generation this season.", 10],
    ["Fall armyworm outbreak in maize field in Karnataka. Severe damage on 30% of plants, spreading fast.", 10],
    ["Desert locust swarm reported in wheat fields of Rajasthan. Need rapid response biopesticide design.", 10],
    ["Colorado potato beetle devastating my potato crop in Himachal Pradesh. Resistance to neonicotinoids suspected.", 8],
    ["Tobacco whitefly infestation in tomato greenhouse in Maharashtra. Mild severity but persistent.", 5],
    ["Peach-potato aphid outbreak in vegetable garden. Organic farm, need bee-safe solution.", 5],
]

# ─────────────────────────────────────────────────────────────────────────────
# Singleton orchestrator
# ─────────────────────────────────────────────────────────────────────────────
_ORCHESTRATOR: BiopesticideOrchestrator | None = None


def get_orchestrator() -> BiopesticideOrchestrator:
    global _ORCHESTRATOR
    if _ORCHESTRATOR is None:
        print("[gradio_app] initializing orchestrator...")
        _ORCHESTRATOR = BiopesticideOrchestrator()
        print(f"[gradio_app] backend = {type(_ORCHESTRATOR.ranker.sirna_model).__name__}, degraded_mode = {_ORCHESTRATOR.degraded_mode}")
    return _ORCHESTRATOR


# ─────────────────────────────────────────────────────────────────────────────
# HTML/CSS helpers
# ─────────────────────────────────────────────────────────────────────────────
CUSTOM_CSS = """
:root {
    --bioai-bg: #f4f5f6;
    --bioai-surface: #ffffff;
    --bioai-card: #ecedee;
    --bioai-accent: #2f86b2;
    --bioai-accent-2: #ba5a6a;
    --bioai-text: #242627;
    --bioai-muted: #71777a;
    --bioai-border: #a1b9c6;
    --bioai-success: #449f63;
    --bioai-warning: #b69045;
    --bioai-error: #964039;
}
.gradio-container { max-width: 1200px !important; }
.bioai-hero {
    background: linear-gradient(135deg, #4a616c 0%, #2f86b2 100%);
    color: white;
    padding: 32px 28px;
    border-radius: 12px;
    margin-bottom: 20px;
}
.bioai-hero h1 {
    font-size: 32px !important;
    font-weight: 800 !important;
    margin: 0 0 8px 0 !important;
    letter-spacing: -0.5px;
}
.bioai-hero p {
    font-size: 14px !important;
    margin: 4px 0 !important;
    opacity: 0.92;
}
.bioai-hero .tagline {
    font-size: 18px !important;
    font-weight: 500 !important;
    margin: 12px 0 4px 0 !important;
}
.bioai-metric-row {
    display: flex;
    gap: 24px;
    margin-top: 20px;
    flex-wrap: wrap;
}
.bioai-metric {
    text-align: center;
}
.bioai-metric .num {
    font-size: 28px;
    font-weight: 800;
    color: white;
}
.bioai-metric .lbl {
    font-size: 11px;
    text-transform: uppercase;
    letter-spacing: 0.5px;
    opacity: 0.85;
    margin-top: 2px;
}
.bioai-card {
    background: var(--bioai-surface);
    border: 1px solid var(--bioai-border);
    border-radius: 8px;
    padding: 16px 20px;
    margin-bottom: 12px;
}
.bioai-status-strip {
    display: flex;
    gap: 8px;
    justify-content: center;
    margin: 8px 0 16px 0;
    flex-wrap: wrap;
}
.bioai-pill {
    padding: 4px 12px;
    border-radius: 12px;
    font-size: 11px;
    font-weight: 600;
    color: white;
}
.bioai-stats-strip {
    display: flex;
    gap: 28px;
    justify-content: center;
    margin: 12px 0 20px 0;
    padding: 14px 0;
    border-top: 1px solid var(--bioai-border);
    border-bottom: 1px solid var(--bioai-border);
    flex-wrap: wrap;
}
.bioai-stat {
    text-align: center;
}
.bioai-stat .num {
    font-size: 22px;
    font-weight: 700;
    color: var(--bioai-accent);
}
.bioai-stat .lbl {
    font-size: 10px;
    color: var(--bioai-muted);
    text-transform: uppercase;
    letter-spacing: 0.5px;
}
.bioai-candidate-card {
    background: var(--bioai-surface);
    border: 1px solid var(--bioai-border);
    border-left: 4px solid var(--bioai-accent);
    border-radius: 6px;
    padding: 14px 18px;
    margin-bottom: 10px;
}
.bioai-candidate-card .header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: 8px;
}
.bioai-candidate-card .rank {
    font-size: 11px;
    font-weight: 700;
    color: var(--bioai-accent);
    text-transform: uppercase;
}
.bioai-candidate-card .seq {
    font-family: monospace;
    font-size: 14px;
    color: var(--bioai-text);
    font-weight: 600;
}
.bioai-candidate-card .metrics {
    display: flex;
    gap: 16px;
    font-size: 12px;
    color: var(--bioai-muted);
}
.bioai-candidate-card .metric-val {
    font-weight: 700;
    color: var(--bioai-text);
}
.bioai-footer {
    text-align: center;
    color: var(--bioai-muted);
    font-size: 11px;
    padding: 16px 0;
    border-top: 1px solid var(--bioai-border);
    margin-top: 24px;
}
"""


def _hero_html() -> str:
    return f"""
    <div class="bioai-hero">
        <h1>Biopesticide-AI</h1>
        <p class="tagline">{TAGLINE}</p>
        <p>{SUBTITLE}</p>
        <div class="bioai-metric-row">
            <div class="bioai-metric"><div class="num">7</div><div class="lbl">Pest species</div></div>
            <div class="bioai-metric"><div class="num">14</div><div class="lbl">Safety panel</div></div>
            <div class="bioai-metric"><div class="num">~4 min</div><div class="lbl">Design loop</div></div>
            <div class="bioai-metric"><div class="num">$0</div><div class="lbl">Cost per design</div></div>
            <div class="bioai-metric"><div class="num">100%</div><div class="lbl">Local compute</div></div>
        </div>
    </div>
    """


def _status_pill(text: str, color: str = "#4a616c") -> str:
    return f'<span class="bioai-pill" style="background:{color};">{text}</span>'


def _backend_status_html(orch: BiopesticideOrchestrator) -> str:
    llm_text = "Ollama Llama 3.2 3B (local)" if not orch.degraded_mode else "Degraded mode (Ollama not running)"
    llm_color = "#449f63" if not orch.degraded_mode else "#b69045"
    model_class = type(orch.ranker.sirna_model).__name__
    model_text = "Caduceus-Ph-1" if model_class == "CaduceusAdapter" else "Dilated CNN (HyenaDNA-inspired)"
    model_color = "#2f86b2" if model_class == "CaduceusAdapter" else "#4a616c"
    device = str(orch.ranker.device)
    return f"""
    <div class="bioai-status-strip">
        {_status_pill('Model: ' + model_text, model_color)}
        {_status_pill('LLM: ' + llm_text, llm_color)}
        {_status_pill('Device: ' + device, '#71777a')}
        {_status_pill('Safety panel: 14 species', '#4c7094')}
        {_status_pill('Pest targets: 7 species', '#ba5a6a')}
    </div>
    """


def _stats_strip(result: dict, elapsed: float) -> str:
    n_tr = result.get("n_transcripts", 0)
    n_pre = result.get("n_precursors", 0)
    n_si = result.get("n_sirnas", 0)
    cost = result.get("total_cost_estimate", 0.0)
    return f"""
    <div class="bioai-stats-strip">
        <div class="bioai-stat"><div class="num">{elapsed:.1f}s</div><div class="lbl">Design loop</div></div>
        <div class="bioai-stat"><div class="num">{n_tr}</div><div class="lbl">Transcripts</div></div>
        <div class="bioai-stat"><div class="num">{n_pre}</div><div class="lbl">Precursors</div></div>
        <div class="bioai-stat"><div class="num">{n_si}</div><div class="lbl">siRNAs scored</div></div>
        <div class="bioai-stat"><div class="num">${cost:.4f}</div><div class="lbl">Est. cost</div></div>
    </div>
    """


def _pest_report_html(pest: dict) -> str:
    if not pest:
        return "<p><i>No pest report parsed.</i></p>"
    species = pest.get("pest_species", pest.get("species", "unknown"))
    crop = pest.get("crop", "unknown")
    severity = pest.get("severity", "unknown")
    location = pest.get("location", "unknown")
    notes = pest.get("notes", "")
    notes_html = f"<tr><td style='padding:4px 18px 4px 0; color:#71777a; font-weight:600; vertical-align:top;'>Notes</td><td style='padding:4px 0; color:#71777a; font-style:italic;'>{notes}</td></tr>" if notes else ""
    return f"""
    <div class="bioai-card">
        <table style="border-collapse:collapse; font-size:13px; width:100%;">
            <tr><td style="padding:4px 18px 4px 0; color:#71777a; font-weight:600; width:140px;">Target species</td><td style="padding:4px 0;"><b>{species}</b></td></tr>
            <tr><td style="padding:4px 18px 4px 0; color:#71777a; font-weight:600;">Crop</td><td style="padding:4px 0;">{crop}</td></tr>
            <tr><td style="padding:4px 18px 4px 0; color:#71777a; font-weight:600;">Severity</td><td style="padding:4px 0;">{severity}</td></tr>
            <tr><td style="padding:4px 18px 4px 0; color:#71777a; font-weight:600;">Location</td><td style="padding:4px 0;">{location}</td></tr>
            {notes_html}
        </table>
    </div>
    """


def _candidate_cards_html(candidates: list) -> str:
    """Render candidates as styled cards instead of a plain table."""
    if not candidates:
        return "<p><i>No candidates generated. Run the design pipeline first.</i></p>"
    cards = []
    for i, c in enumerate(candidates, 1):
        seq = c.get("sirna_seq", "")
        eff = c.get("efficacy", 0)
        ot = c.get("offtarget_max", 0)
        hl = c.get("half_life_hours", 0)
        score = c.get("final_score", 0)
        hl_days = hl / 24
        # Risk tier color for the left border
        if score > 0.3:
            border_color = "#449f63"  # success
        elif score > 0.15:
            border_color = "#b69045"  # warning
        else:
            border_color = "#964039"  # error
        cards.append(f"""
        <div class="bioai-candidate-card" style="border-left-color:{border_color};">
            <div class="header">
                <span class="rank">#{i}</span>
                <span class="seq">{seq}</span>
            </div>
            <div class="metrics">
                <span>Efficacy: <span class="metric-val">{eff:.3f}</span></span>
                <span>Off-target max: <span class="metric-val">{ot:.3f}</span></span>
                <span>Half-life: <span class="metric-val">{hl:.1f}h ({hl_days:.1f}d)</span></span>
                <span>Final score: <span class="metric-val">{score:.3f}</span></span>
            </div>
        </div>
        """)
    return "".join(cards)


def _candidates_to_dataframe(candidates: list) -> list:
    rows = []
    for i, c in enumerate(candidates, 1):
        rows.append([
            i,
            c.get("sirna_seq", ""),
            f"{c.get('efficacy', 0):.3f}",
            f"{c.get('offtarget_max', 0):.3f}",
            f"{c.get('half_life_hours', 0):.1f}h",
            f"{c.get('final_score', 0):.3f}",
        ])
    return rows


def _safety_cards_md(result: dict) -> str:
    sc = result.get("safety_cards", "")
    if isinstance(sc, list):
        parts = []
        for i, c in enumerate(sc, 1):
            if isinstance(c, dict):
                parts.append(f"### Candidate #{i}: `{c.get('sirna_seq', '')}`\n\n{c.get('card_markdown', '')}")
            else:
                parts.append(str(c))
        return "\n\n---\n\n".join(parts) if parts else "_(no safety cards generated)_"
    return sc or "_(no safety cards generated)_"


def _export_csv(candidates: list) -> str:
    """Generate CSV string for download."""
    if not candidates:
        return ""
    output = io.StringIO()
    writer = csv.writer(output)
    writer.writerow(["rank", "sirna_seq", "efficacy", "offtarget_max", "half_life_hours", "final_score"])
    for i, c in enumerate(candidates, 1):
        writer.writerow([
            i,
            c.get("sirna_seq", ""),
            f"{c.get('efficacy', 0):.4f}",
            f"{c.get('offtarget_max', 0):.4f}",
            f"{c.get('half_life_hours', 0):.2f}",
            f"{c.get('final_score', 0):.4f}",
        ])
    return output.getvalue()


def _export_json(result: dict) -> str:
    """Generate JSON string for download."""
    slim = {k: v for k, v in result.items() if k not in ("safety_cards", "regulatory_memo")}
    if "candidates" in slim and isinstance(slim["candidates"], list):
        slim["candidates"] = slim["candidates"][:10]
    return json.dumps(slim, indent=2, default=str)


# ─────────────────────────────────────────────────────────────────────────────
# Main design handler
# ─────────────────────────────────────────────────────────────────────────────
def design_handler(user_text: str, top_k: int):
    """Run the design pipeline and return all UI outputs."""
    if not user_text or not user_text.strip():
        empty_status = "<p style='color:#964039;'><b>Please describe your pest problem above.</b></p>"
        return (
            empty_status,
            gr.update(value=[]),
            "<p><i>No pest report parsed.</i></p>",
            "",
            None, None, None,  # charts
            "",
            "_(no safety cards generated)_",
            "_(no regulatory memo generated)_",
            "",
            "",
        )

    orch = get_orchestrator()
    t0 = time.time()
    try:
        result = orch.design(user_text, top_k=int(top_k))
    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        err = f"<p style='color:#964039;'><b>Pipeline error:</b> {type(e).__name__}: {e}</p><pre style='font-size:10px;'>{tb}</pre>"
        return (err, gr.update(value=[]), "<p><i>No pest report parsed.</i></p>", "", None, None, None, "", "_(no safety cards generated)_", "_(no regulatory memo generated)_", "", "")
    elapsed = time.time() - t0

    candidates = result.get("candidates", [])
    status_html = _backend_status_html(orch) + _stats_strip(result, elapsed)
    candidates_rows = _candidates_to_dataframe(candidates)
    candidates_html = _candidate_cards_html(candidates)
    pest_html = _pest_report_html(result.get("pest_report", {}))
    safety_md = _safety_cards_md(result)
    memo_md = result.get("regulatory_memo", "_(no regulatory memo generated)_")
    csv_str = _export_csv(candidates)
    json_str = _export_json(result)

    # Generate charts
    efficacy_chart = efficacy_bar_chart(candidates) if candidates else None
    offtarget_chart = offtarget_heatmap(candidates, SAFETY_SPECIES) if candidates else None
    halflife_chart_path = halflife_chart(candidates) if candidates else None

    return (
        status_html,
        candidates_rows,
        candidates_html,
        pest_html,
        "",
        efficacy_chart,
        offtarget_chart,
        halflife_chart_path,
        safety_md,
        memo_md,
        csv_str,
        json_str,
    )


# ─────────────────────────────────────────────────────────────────────────────
# Gradio Blocks UI
# ─────────────────────────────────────────────────────────────────────────────
def build_ui() -> gr.Blocks:
    demo = gr.Blocks(title="Biopesticide-AI")

    with demo:
        # ─── Hero ─────────────────────────────────────────────────────────
        gr.HTML(_hero_html())

        # ─── Initial status ───────────────────────────────────────────────
        try:
            orch = get_orchestrator()
            initial_status = _backend_status_html(orch) + "<div style='text-align:center; color:#71777a; font-size:12px; padding:10px 0;'>Click <b>Design dsRNA candidates</b> to run the pipeline.</div>"
        except Exception as e:
            initial_status = f"<p style='color:#964039;'>Failed to initialize orchestrator: {e}</p>"

        status_box = gr.HTML(value=initial_status, label="Pipeline status")

        # ─── Main tabs ────────────────────────────────────────────────────
        with gr.Tabs():
            # ── Tab 1: Design ─────────────────────────────────────────────
            with gr.Tab("Design", id=0):
                gr.Markdown("### Describe your pest problem in plain English")
                user_text = gr.Textbox(
                    label="Pest report",
                    placeholder="e.g. 'Brown planthopper infestation in my rice paddy near Coimbatore, Tamil Nadu. Severity moderate, second generation this season.'",
                    lines=4,
                    value=EXAMPLES[0][0],
                )
                with gr.Accordion("Advanced settings", open=False):
                    top_k = gr.Slider(minimum=1, maximum=20, value=10, step=1, label="Top-K candidates to return")

                run_btn = gr.Button("Design dsRNA candidates", variant="primary", size="lg")

                gr.Examples(
                    examples=EXAMPLES,
                    inputs=[user_text, top_k],
                    label="Try one of these preset pest reports",
                )

                gr.Markdown("---")
                gr.Markdown("### Parsed pest report")
                pest_html = gr.HTML(value="<p style='color:#71777a;'><i>Run the pipeline to see the parsed pest report.</i></p>")

                gr.Markdown("### Top candidates")
                candidates_html = gr.HTML(value="<p style='color:#71777a;'><i>Run the pipeline to see ranked candidates.</i></p>")

                # Hidden dataframe for CSV export compatibility
                candidates_table = gr.Dataframe(
                    visible=False,
                    headers=["Rank", "siRNA", "Efficacy", "Off-target", "Half-life", "Score"],
                    value=[],
                )

            # ── Tab 2: Analytics ──────────────────────────────────────────
            with gr.Tab("Analytics", id=1):
                gr.Markdown("### Efficacy scores")
                efficacy_img = gr.Image(label="", show_label=False, height=350)
                gr.Markdown("### Off-target risk heatmap (candidates x 14 safety species)")
                offtarget_img = gr.Image(label="", show_label=False, height=400)
                gr.Markdown("### Environmental fate (predicted half-life)")
                halflife_img = gr.Image(label="", show_label=False, height=350)

            # ── Tab 3: Safety ─────────────────────────────────────────────
            with gr.Tab("Safety cards", id=2):
                gr.Markdown("### Per-candidate safety cards (generated by local Llama 3.2 3B)")
                safety_md = gr.Markdown(value="<p><i>Run the pipeline to see safety cards for the top candidates.</i></p>")

            # ── Tab 4: Regulatory ─────────────────────────────────────────
            with gr.Tab("Regulatory memo", id=3):
                gr.Markdown("### EPA-style regulatory memo (generated by local Llama 3.2 3B)")
                memo_md = gr.Markdown(value="<p><i>Run the pipeline to see the regulatory memo.</i></p>")

            # ── Tab 5: Export ─────────────────────────────────────────────
            with gr.Tab("Export", id=4):
                gr.Markdown("### Download design results")
                gr.Markdown("Export the top candidates as CSV or the full design result as JSON for downstream analysis.")
                csv_text = gr.Textbox(label="CSV (copy below or use the download button)", lines=10, interactive=False)
                csv_btn = gr.DownloadButton("Download CSV", value=None)
                json_text = gr.Textbox(label="JSON (copy below or use the download button)", lines=15, interactive=False)
                json_btn = gr.DownloadButton("Download JSON", value=None)

            # ── Tab 6: About ──────────────────────────────────────────────
            with gr.Tab("About", id=5):
                gr.Markdown("""
                ### About Biopesticide-AI

                **Biopesticide-AI** is an end-to-end pipeline for designing dsRNA biopesticides against agricultural pests. It compresses the traditional 3-6 month wet-lab design loop into a 4-minute computational pipeline that any farmer, agronomist, or cooperative can run from a laptop.

                **Pipeline:**
                1. Farmer describes pest problem in plain English
                2. Local Ollama Llama 3.2 3B parses the report into a structured design spec
                3. PyTorch backend tiles pest transcripts into 200-nt dsRNA precursors
                4. Dicer-style dicing produces 21-nt siRNAs
                5. Dilated CNN (HyenaDNA-inspired) scores each siRNA for efficacy
                6. K-mer index checks off-target risk against 14 non-target species
                7. Physics-Informed Neural Network predicts environmental half-life
                8. Learned ranker combines all scores into a final candidate ranking
                9. Llama 3.2 3B generates safety cards + EPA-style regulatory memo

                **14-species safety panel** covers pollinators (honeybee, bumblebee, leafcutter bee), beneficial predators (ladybug, lacewing), soil invertebrates (earthworm), aquatic organisms (water flea, zebrafish), livestock (cattle, zebu, chicken, sheep, pig), and human safety.

                **7 pest targets** include brown planthopper (rice), fall armyworm (maize), desert locust (wheat), striped stem borer (rice), peach-potato aphid (vegetables), Colorado potato beetle (potato), and tobacco whitefly (tomato).

                **Cost: $0 per design.** All compute is local. No cloud API spend.

                Built for the AMD Developer Hackathon Unicorn Track. MIT licensed.
                """)

        # ─── Footer ───────────────────────────────────────────────────────
        gr.HTML(
            "<div class='bioai-footer'>"
            "Built for the AMD Developer Hackathon Unicorn Track. "
            "Backend: PyTorch + Caduceus (with CNN fallback). "
            "LLM: Ollama Llama 3.2 3B running locally. "
            "14-species safety panel. 7 pest targets. "
            "Containerized via Docker. MIT licensed."
            "</div>"
        )

        # ─── Wire up ──────────────────────────────────────────────────────
        run_btn.click(
            design_handler,
            inputs=[user_text, top_k],
            outputs=[
                status_box,
                candidates_table,
                candidates_html,
                pest_html,
                status_box,  # update status after run (same component)
                efficacy_img,
                offtarget_img,
                halflife_img,
                safety_md,
                memo_md,
                csv_text,
                json_text,
            ],
        )

    return demo


# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Biopesticide-AI Gradio UI v2")
    parser.add_argument("--host", default="0.0.0.0", help="bind host (default 0.0.0.0)")
    parser.add_argument("--port", type=int, default=7860, help="bind port (default 7860)")
    parser.add_argument("--share", action="store_true", help="create a public share link")
    parser.add_argument("--max-threads", type=int, default=4, help="max concurrent requests")
    args = parser.parse_args()

    print("[gradio_app] pre-initializing orchestrator...")
    get_orchestrator()

    demo = build_ui()
    print(f"[gradio_app] launching on http://{args.host}:{args.port}")
    demo.launch(
        server_name=args.host,
        server_port=args.port,
        share=args.share,
        max_threads=args.max_threads,
        show_error=True,
        css=CUSTOM_CSS,
    )


if __name__ == "__main__":
    main()