#!/usr/bin/env python
from __future__ import annotations

import json
import os
from functools import lru_cache
from pathlib import Path
from typing import Any

from assaylens.agent.model_predictor import predictor_from_artifacts
from assaylens.agent.next_assay import recommend_next_assay


DEFAULT_ARTIFACT_REPO = "lighteternal/assaylens-artifacts"


DEMO_CASES = [
    {
        "label": "7TM agonist signal needing pathway contrast",
        "request": {
            "compound": {"compound_id": "demo_cmpd_7tm_agonist", "name": "demo agonist", "smiles": "CCN(CC)CCO"},
            "target": {
                "target_id": "demo_tgt_7tm",
                "gene": "ADRB2",
                "target_class": "7TM",
                "uniprot_id": "P07550",
            },
            "observed_result": {
                "mode": "Agonist",
                "assay__mechanism": "G Protein Activation",
                "assay__detailed_mechanism": "cAMP",
                "assay__technology": "Luminescence",
                "outcome_is_active": True,
                "outcome_potency_pxc50": 6.8,
                "outcome_max_activity": 82.0,
                "is_quantified": True,
                "frequency_flag": False,
                "viability_flag": False,
            },
            "available_assay_frames": [
                {"mode": "Agonist", "assay__mechanism": "Barr2 Recruitment", "assay__technology": "Unknown"},
                {"mode": "Antagonist", "assay__mechanism": "G Protein Activation", "assay__technology": "Unknown"},
                {"mode": "Binding", "assay__mechanism": "Competition Binding", "assay__technology": "TR-FRET"},
            ],
        },
    },
    {
        "label": "Kinase binding signal needing orthogonal confirmation",
        "request": {
            "compound": {"compound_id": "demo_cmpd_kinase", "name": "demo kinase ligand", "smiles": "CC1=NC=CC=C1"},
            "target": {
                "target_id": "demo_tgt_kinase",
                "gene": "ABL1",
                "target_class": "Kinase",
                "uniprot_id": "P00519",
            },
            "observed_result": {
                "mode": "Binding",
                "assay__mechanism": "Competition Binding",
                "assay__detailed_mechanism": "competition",
                "assay__technology": "TR-FRET",
                "outcome_is_active": True,
                "outcome_potency_pxc50": 7.1,
                "outcome_max_activity": 65.0,
                "is_quantified": True,
                "frequency_flag": True,
                "viability_flag": False,
            },
            "available_assay_frames": [
                {"mode": "Inhibition", "assay__mechanism": "Enzyme Activity", "assay__technology": "Unknown"},
                {"mode": "Binding", "assay__mechanism": "Thermal Shift", "assay__technology": "Unknown"},
                {"mode": "Binding", "assay__mechanism": "Competition Binding", "assay__technology": "TR-FRET"},
            ],
        },
    },
]


def default_request() -> dict[str, Any]:
    path = Path("examples/tiny_observations.json")
    if not path.exists():
        return {}
    payload = json.loads(path.read_text(encoding="utf-8"))
    return payload[0] if isinstance(payload, list) and payload else payload


def demo_cases() -> list[dict[str, Any]]:
    cases = [{"label": "Viability-confounded 7TM antagonist", "request": default_request()}]
    for case in DEMO_CASES:
        request = case["request"]() if callable(case["request"]) else case["request"]
        cases.append({"label": case["label"], "request": request})
    return cases


def default_request_text() -> str:
    return json.dumps(default_request(), indent=2)


def default_candidate_frames_text() -> str:
    return json.dumps(default_request().get("available_assay_frames", []), indent=2)


def case_by_label(label: str) -> dict[str, Any]:
    for case in demo_cases():
        if case["label"] == label:
            return case["request"]
    return default_request()


def fields_from_request(request: dict[str, Any]) -> tuple[Any, ...]:
    compound = request.get("compound", {})
    target = request.get("target", {})
    observed = request.get("observed_result", {})
    return (
        compound.get("compound_id", ""),
        compound.get("name", ""),
        compound.get("smiles", ""),
        target.get("target_id", ""),
        target.get("gene", ""),
        target.get("target_class", ""),
        target.get("uniprot_id", ""),
        observed.get("mode", ""),
        observed.get("assay__mechanism", ""),
        observed.get("assay__detailed_mechanism", ""),
        observed.get("assay__technology", ""),
        bool(observed.get("outcome_is_active")),
        bool(observed.get("is_quantified")),
        bool(observed.get("frequency_flag")),
        bool(observed.get("viability_flag")),
        json.dumps(request.get("available_assay_frames", []), indent=2),
        json.dumps(request, indent=2),
    )


def load_case(label: str) -> tuple[Any, ...]:
    return fields_from_request(case_by_label(label))


@lru_cache(maxsize=4)
def cached_predictor(artifact_repo: str, model_artifact: str):
    return predictor_from_artifacts(artifact_repo, model_artifact)


@lru_cache(maxsize=4)
def cached_frame_coverage(artifact_repo: str) -> dict[str, int] | None:
    try:
        import pandas as pd
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(artifact_repo, "data_profiles/frame_counts.csv", repo_type="dataset")
        frame = pd.read_csv(path)
        return {str(row["assay_frame_coarse"]): int(row["n_rows"]) for _, row in frame.iterrows()}
    except Exception:  # noqa: BLE001
        return None


def run_recommendation(
    request_text: str,
    artifact_repo: str,
    model_artifact: str,
    include_llm: bool,
) -> tuple[str, list[dict[str, Any]], str, str, str]:
    request = json.loads(request_text)
    return recommend_request(request, artifact_repo, model_artifact, include_llm)


def run_form_recommendation(
    compound_id: str,
    compound_name: str,
    smiles: str,
    target_id: str,
    gene: str,
    target_class: str,
    uniprot_id: str,
    mode: str,
    mechanism: str,
    detailed_mechanism: str,
    technology: str,
    active: bool,
    quantified: bool,
    frequency_flag: bool,
    viability_flag: bool,
    candidate_frames_text: str,
    artifact_repo: str,
    model_artifact: str,
    include_llm: bool,
) -> tuple[str, list[dict[str, Any]], str, str, str]:
    request = {
        "compound": {
            "compound_id": compound_id,
            "name": compound_name,
            "smiles": smiles,
        },
        "target": {
            "target_id": target_id,
            "gene": gene,
            "target_class": target_class,
            "uniprot_id": uniprot_id,
        },
        "observed_result": {
            "mode": mode,
            "assay__mechanism": mechanism,
            "assay__detailed_mechanism": detailed_mechanism,
            "assay__technology": technology,
            "outcome_is_active": active,
            "is_quantified": quantified,
            "frequency_flag": frequency_flag,
            "viability_flag": viability_flag,
        },
        "available_assay_frames": json.loads(candidate_frames_text),
    }
    return recommend_request(request, artifact_repo, model_artifact, include_llm)


def recommend_request(
    request: dict[str, Any],
    artifact_repo: str,
    model_artifact: str,
    include_llm: bool,
) -> tuple[str, list[dict[str, Any]], str, str, str]:
    predictor = cached_predictor(artifact_repo, model_artifact)
    frame_coverage = cached_frame_coverage(artifact_repo)
    result = recommend_next_assay(
        request,
        assaylens_mode="COUNTERFACTUAL_READY",
        frame_coverage=frame_coverage,
        predictor=predictor,
        include_llm_explanation=include_llm,
    )
    recommendation = result.get("recommended_next_assay", {})
    frame = recommendation.get("frame", {})
    summary = "\n".join(
        [
            f"Recommended frame: {frame.get('mode', 'unknown')} / {frame.get('assay__mechanism', 'unknown')} / {frame.get('assay__technology', 'unknown')}",
            f"Confidence: {recommendation.get('confidence', 'unknown')}",
            f"Coverage: {recommendation.get('coverage_status', 'unknown')} ({recommendation.get('coverage_count', 'n/a')})",
            f"Model source: {recommendation.get('model_source', 'unknown')}",
            f"Why: {recommendation.get('reason', 'unknown')}",
        ]
    )
    ranked = []
    for idx, candidate in enumerate(result.get("counterfactual_predictions", []), start=1):
        candidate_frame = candidate.get("candidate_frame", {})
        ranked.append(
            {
                "rank": idx,
                "mode": candidate_frame.get("mode"),
                "mechanism": candidate_frame.get("assay__mechanism"),
                "technology": candidate_frame.get("assay__technology"),
                "active": round(float(candidate.get("predicted_active_probability", 0.0)), 4),
                "quantified": round(float(candidate.get("predicted_quantified_probability", 0.0)), 4),
                "viability": round(float(candidate.get("predicted_viability_confounding_probability", 0.0)), 4),
                "frequency": round(float(candidate.get("predicted_frequency_artifact_probability", 0.0)), 4),
                "uncertainty": round(float(candidate.get("uncertainty_score", 0.0)), 4),
                "distance": round(float(candidate.get("frame_distance", 0.0)), 4),
                "artifact_score": round(float(candidate.get("artifact_disambiguation_score", 0.0)), 4),
                "coverage": candidate.get("coverage_count"),
                "coverage_status": candidate.get("coverage_status"),
                "total_score": round(float(candidate.get("total_score", candidate.get("score", 0.0))), 4),
                "confidence": candidate.get("recommendation_confidence"),
                "exploratory": candidate.get("exploratory"),
            }
        )
    warnings = "\n".join(f"- {warning}" for warning in result.get("warnings", []))
    observed = result.get("observed_signal_interpretation", {})
    evidence = "\n".join(
        [
            f"Observed active probability: {float(observed.get('active_probability_under_observed_frame', 0.0)):.4f}",
            f"Viability confounding probability: {float(observed.get('viability_confounding_probability', 0.0)):.4f}",
            f"Frequency artifact probability: {float(observed.get('frequency_artifact_probability', 0.0)):.4f}",
            f"Observed-frame model source: {observed.get('model_source', 'unknown')}",
        ]
    )
    return (
        summary,
        ranked,
        json.dumps(result.get("agent_trace", {}), indent=2),
        result.get("llm_explanation") or warnings,
        evidence,
    )


def build_app():
    import gradio as gr

    sample = default_request()
    compound = sample.get("compound", {})
    target = sample.get("target", {})
    observed = sample.get("observed_result", {})
    css = """
    .gradio-container {
      background: #f7f7f2;
      color: #1e2320;
      font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
    }
    #assaylens-title {
      border-left: 6px solid #276357;
      padding-left: 16px;
      letter-spacing: 0;
    }
    .assaylens-panel {
      border: 1px solid #c7d0ca;
      background: #ffffff;
      border-radius: 8px;
    }
    button.primary {
      background: #276357 !important;
      color: #ffffff !important;
      border: 1px solid #1e2320 !important;
    }
    """
    with gr.Blocks(css=css, title="AssayLens") as app:
        gr.Markdown(
            "# AssayLens\nRetrospective assay-frame evidence and next-measurement recommendation.",
            elem_id="assaylens-title",
        )
        with gr.Row(equal_height=True):
            case_picker = gr.Dropdown(
                choices=[case["label"] for case in demo_cases()],
                value=demo_cases()[0]["label"],
                label="Curated case",
            )
            load_case_button = gr.Button("Load Case")
            run_form = gr.Button("Run Recommendation", variant="primary")
        with gr.Accordion("Advanced artifact controls", open=False):
            with gr.Row(equal_height=True):
                artifact_repo = gr.Textbox(value=DEFAULT_ARTIFACT_REPO, label="Artifact repo")
                model_artifact = gr.Dropdown(
                    choices=[
                        "models/b3_ligand_target_assay.pkl",
                        "models/b4_ligand_target_assay_goa.pkl",
                    ],
                    value="models/b3_ligand_target_assay.pkl",
                    label="Model artifact",
                )
                include_llm = gr.Checkbox(
                    value=False,
                    label="Include LLM explanation when OPENAI_API_KEY is set",
                )
        with gr.Row():
            with gr.Column(scale=5, elem_classes=["assaylens-panel"]):
                with gr.Tabs():
                    with gr.Tab("Structured input"):
                        with gr.Row():
                            compound_id = gr.Textbox(value=compound.get("compound_id", ""), label="Compound ID")
                            compound_name = gr.Textbox(value=compound.get("name", ""), label="Compound name")
                        smiles = gr.Textbox(value=compound.get("smiles", ""), label="SMILES")
                        with gr.Row():
                            target_id = gr.Textbox(value=target.get("target_id", ""), label="Target ID")
                            gene = gr.Textbox(value=target.get("gene", ""), label="Gene")
                        with gr.Row():
                            target_class = gr.Textbox(value=target.get("target_class", ""), label="Target class")
                            uniprot_id = gr.Textbox(value=target.get("uniprot_id", ""), label="UniProt ID")
                        with gr.Row():
                            mode = gr.Textbox(value=observed.get("mode", ""), label="Observed mode")
                            mechanism = gr.Textbox(value=observed.get("assay__mechanism", ""), label="Mechanism")
                        with gr.Row():
                            detailed_mechanism = gr.Textbox(
                                value=observed.get("assay__detailed_mechanism", ""),
                                label="Detailed mechanism",
                            )
                            technology = gr.Textbox(value=observed.get("assay__technology", ""), label="Technology")
                        with gr.Row():
                            active = gr.Checkbox(value=bool(observed.get("outcome_is_active")), label="Active")
                            quantified = gr.Checkbox(value=bool(observed.get("is_quantified")), label="Quantified")
                            frequency_flag = gr.Checkbox(
                                value=bool(observed.get("frequency_flag")),
                                label="Frequency artifact",
                            )
                            viability_flag = gr.Checkbox(
                                value=bool(observed.get("viability_flag")),
                                label="Viability confound",
                            )
                        candidate_frames = gr.Code(
                            value=default_candidate_frames_text(),
                            language="json",
                            label="Candidate assay frames",
                        )
                    with gr.Tab("JSON input"):
                        request = gr.Code(value=default_request_text(), language="json", label="Observed assay request")
                        run_json = gr.Button("Run JSON Recommendation", variant="primary")
            with gr.Column(scale=4, elem_classes=["assaylens-panel"]):
                summary = gr.Textbox(label="Recommendation", lines=4)
                evidence = gr.Textbox(label="Observed evidence", lines=5)
                ranked = gr.Dataframe(label="Ranked candidate frames")
                explanation = gr.Textbox(label="Warnings / LLM explanation", lines=8)
        trace = gr.Code(language="json", label="Agent trace")
        load_case_button.click(
            load_case,
            inputs=[case_picker],
            outputs=[
                compound_id,
                compound_name,
                smiles,
                target_id,
                gene,
                target_class,
                uniprot_id,
                mode,
                mechanism,
                detailed_mechanism,
                technology,
                active,
                quantified,
                frequency_flag,
                viability_flag,
                candidate_frames,
                request,
            ],
        )
        run_form.click(
            run_form_recommendation,
            inputs=[
                compound_id,
                compound_name,
                smiles,
                target_id,
                gene,
                target_class,
                uniprot_id,
                mode,
                mechanism,
                detailed_mechanism,
                technology,
                active,
                quantified,
                frequency_flag,
                viability_flag,
                candidate_frames,
                artifact_repo,
                model_artifact,
                include_llm,
            ],
            outputs=[summary, ranked, trace, explanation, evidence],
        )
        run_json.click(
            run_recommendation,
            inputs=[request, artifact_repo, model_artifact, include_llm],
            outputs=[summary, ranked, trace, explanation, evidence],
        )
    return app


def main() -> None:
    app = build_app()
    app.launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1"))


if __name__ == "__main__":
    main()