training/literature_benchmark.py

"""Literature-grounded experiment benchmark utilities.

This module lets the environment run a paper-backed experiment plan, then
compare the resulting simulated findings against curated expected findings
from the literature.
"""

from __future__ import annotations

import argparse
import json
import re
from dataclasses import asdict, dataclass, field
from importlib.metadata import PackageNotFoundError, version
from typing import Any, Dict, List, Optional, Sequence

from models import (
    ActionType,
    ConclusionClaim,
    ExperimentAction,
    ExperimentObservation,
    OutputType,
    TaskSpec,
)
from server.hackathon_environment import BioExperimentEnvironment
from server.tasks.scenarios import SCENARIO_LIBRARY, Scenario

TOKEN_RE = re.compile(r"[A-Za-z0-9_+\-]+")
STOPWORDS = {
    "a",
    "an",
    "and",
    "as",
    "by",
    "for",
    "from",
    "in",
    "into",
    "of",
    "on",
    "or",
    "the",
    "to",
    "using",
    "with",
}

BIO_LIBRARY_DISTRIBUTIONS = {
    "scanpy": "scanpy",
    "gseapy": "gseapy",
    "biopython": "biopython",
}


@dataclass
class PaperBenchmarkResult:
    scenario_name: str
    problem_statement: str
    matched_papers: List[str]
    bio_library_versions: Dict[str, Optional[str]]
    matched_findings: List[str] = field(default_factory=list)
    missed_findings: List[str] = field(default_factory=list)
    discovered_markers: List[str] = field(default_factory=list)
    candidate_mechanisms: List[str] = field(default_factory=list)
    conclusions: List[str] = field(default_factory=list)
    final_reward: float = 0.0
    total_steps: int = 0

    @property
    def match_ratio(self) -> float:
        total = len(self.matched_findings) + len(self.missed_findings)
        return len(self.matched_findings) / max(total, 1)

    def to_dict(self) -> Dict[str, Any]:
        data = asdict(self)
        data["match_ratio"] = self.match_ratio
        return data


def detect_bio_library_versions() -> Dict[str, Optional[str]]:
    versions: Dict[str, Optional[str]] = {}
    for name, dist_name in BIO_LIBRARY_DISTRIBUTIONS.items():
        try:
            versions[name] = version(dist_name)
        except PackageNotFoundError:
            versions[name] = None
    return versions


def select_literature_scenario(problem_statement: str) -> Scenario:
    """Pick the closest literature-backed scenario for a prompt."""

    prompt_tokens = set(_tokenize(problem_statement))
    best_score = -1
    best_scenario: Optional[Scenario] = None

    for scenario in SCENARIO_LIBRARY:
        if not scenario.task.paper_references:
            continue
        corpus = [
            scenario.task.problem_statement,
            *(ref.title for ref in scenario.task.paper_references),
            *(finding.finding for finding in scenario.task.expected_findings),
            scenario.task.tissue,
            scenario.task.modality,
            *scenario.task.conditions,
        ]
        score = len(prompt_tokens & set(_tokenize(" ".join(corpus))))
        if scenario.task.problem_statement.lower() in problem_statement.lower():
            score += 4
        if score > best_score:
            best_score = score
            best_scenario = scenario

    if best_scenario is None:
        raise ValueError("No literature-backed scenarios are available.")
    return best_scenario


def run_paper_benchmark(
    *,
    problem_statement: str,
    scenario_name: Optional[str] = None,
    domain_randomise: bool = False,
) -> PaperBenchmarkResult:
    """Run a literature-backed episode and compare outputs to paper results."""

    scenario = _resolve_scenario(problem_statement, scenario_name)
    env = BioExperimentEnvironment(
        scenario_name=scenario.name,
        domain_randomise=domain_randomise,
    )
    obs = env.reset()

    for action in build_paper_aligned_actions(obs.task):
        obs = env.step(action)

    claims = infer_conclusion_claims(obs)
    obs = env.step(
        ExperimentAction(
            action_type=ActionType.SYNTHESIZE_CONCLUSION,
            parameters={"claims": [claim.model_dump() for claim in claims]},
            justification=(
                "Summarize the simulated experimental evidence and compare it "
                "with the paper-backed expected findings."
            ),
            confidence=0.8,
            tool_call_spec=_tool_context(
                obs.task,
                libraries=["biopython"],
            ),
        )
    )

    matched, missed = compare_expected_findings(obs.task, obs)
    return PaperBenchmarkResult(
        scenario_name=scenario.name,
        problem_statement=obs.task.problem_statement,
        matched_papers=[ref.title for ref in obs.task.paper_references],
        bio_library_versions=detect_bio_library_versions(),
        matched_findings=matched,
        missed_findings=missed,
        discovered_markers=list(obs.discovered_markers),
        candidate_mechanisms=list(obs.candidate_mechanisms),
        conclusions=[c.claim for c in obs.conclusions],
        final_reward=float(obs.metadata.get("cumulative_reward", 0.0)),
        total_steps=obs.step_index,
    )


def build_paper_aligned_actions(task: TaskSpec) -> List[ExperimentAction]:
    """Construct a pragmatic analysis plan aligned to the task modality."""

    actions: List[ExperimentAction] = [
        ExperimentAction(
            action_type=ActionType.COLLECT_SAMPLE,
            parameters={"n_samples": 8},
            justification="Collect enough samples to support downstream analysis.",
            confidence=0.75,
            tool_call_spec=_tool_context(task, libraries=["biopython"]),
        ),
        ExperimentAction(
            action_type=ActionType.PREPARE_LIBRARY,
            method="10x_chromium",
            justification="Use a standard single-cell library prep workflow.",
            confidence=0.8,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.SEQUENCE_CELLS,
            method="NovaSeq",
            justification="Generate sufficient single-cell read depth.",
            confidence=0.8,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.RUN_QC,
            method="scanpy.pp.calculate_qc_metrics",
            justification="Check technical quality before downstream inference.",
            confidence=0.85,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.FILTER_DATA,
            method="scanpy.pp.filter_cells",
            justification="Remove low-quality cells and reduce technical noise.",
            confidence=0.85,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.NORMALIZE_DATA,
            method="scanpy.pp.normalize_total",
            justification="Normalize expression to prepare comparable profiles.",
            confidence=0.85,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.CLUSTER_CELLS,
            method="scanpy.tl.leiden",
            justification="Resolve cell states before focused interpretation.",
            confidence=0.8,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
    ]

    categories = {finding.category for finding in task.expected_findings}
    if "trajectory" in categories:
        actions.extend([
            ExperimentAction(
                action_type=ActionType.TRAJECTORY_ANALYSIS,
                method="scanpy.tl.dpt",
                justification="Recover pseudotime structure and lineage branches.",
                confidence=0.8,
                tool_call_spec=_tool_context(task, libraries=["scanpy"]),
            ),
            ExperimentAction(
                action_type=ActionType.REGULATORY_NETWORK_INFERENCE,
                method="pySCENIC",
                justification="Infer branch-associated regulators from the trajectory.",
                confidence=0.75,
                tool_call_spec=_tool_context(task, libraries=["scanpy"]),
            ),
            ExperimentAction(
                action_type=ActionType.MARKER_SELECTION,
                method="scanpy.tl.rank_genes_groups",
                justification="Summarize lineage markers and branch-state genes.",
                confidence=0.75,
                tool_call_spec=_tool_context(task, libraries=["scanpy"]),
            ),
        ])
        return actions

    actions.extend([
        ExperimentAction(
            action_type=ActionType.DIFFERENTIAL_EXPRESSION,
            method="scanpy.tl.rank_genes_groups",
            parameters={"comparison": _default_comparison_name(task)},
            justification="Identify genes associated with the focal phenotype.",
            confidence=0.85,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.PATHWAY_ENRICHMENT,
            method="gseapy.prerank",
            justification="Translate DE hits into pathway-level interpretation.",
            confidence=0.8,
            tool_call_spec=_tool_context(task, libraries=["gseapy"]),
        ),
        ExperimentAction(
            action_type=ActionType.MARKER_SELECTION,
            method="scanpy.tl.rank_genes_groups",
            justification="Nominate candidate markers for follow-up validation.",
            confidence=0.8,
            tool_call_spec=_tool_context(task, libraries=["scanpy"]),
        ),
        ExperimentAction(
            action_type=ActionType.VALIDATE_MARKER,
            method="immunofluorescence",
            parameters={"marker": _preferred_marker(task)},
            justification="Check whether the leading marker reproduces in validation.",
            confidence=0.75,
            tool_call_spec=_tool_context(task, libraries=["biopython"]),
        ),
    ])
    return actions


def infer_conclusion_claims(obs: ExperimentObservation) -> List[ConclusionClaim]:
    """Turn accumulated evidence into concise, paper-comparable claims."""

    markers = set(obs.discovered_markers)
    mechanisms = set(obs.candidate_mechanisms)
    network_regulators = set(_extract_network_regulators(obs))
    trajectory_output = _latest_output_data(obs, OutputType.TRAJECTORY_RESULT)

    claims: List[ConclusionClaim] = []

    if "SPP1" in markers:
        claims.append(ConclusionClaim(
            claim="SPP1-positive macrophages are enriched in IPF fibrotic tissue.",
            confidence=0.84,
            claim_type="marker",
            evidence_steps=_evidence_steps(obs, {
                OutputType.DE_RESULT,
                OutputType.MARKER_RESULT,
                OutputType.VALIDATION_RESULT,
            }),
        ))
    if {"SPP1", "MERTK"} <= markers:
        claims.append(ConclusionClaim(
            claim="MERTK co-occurs with the SPP1-positive profibrotic macrophage state.",
            confidence=0.8,
            claim_type="marker",
            evidence_steps=_evidence_steps(obs, {
                OutputType.DE_RESULT,
                OutputType.MARKER_RESULT,
            }),
        ))
    if "extracellular_matrix_organisation" in mechanisms:
        claims.append(ConclusionClaim(
            claim=(
                "Extracellular matrix organization is a dominant fibrotic "
                "program in the IPF samples."
            ),
            confidence=0.78,
            claim_type="pathway",
            evidence_steps=_evidence_steps(obs, {OutputType.PATHWAY_RESULT}),
        ))

    if trajectory_output.get("branching_detected"):
        claims.append(ConclusionClaim(
            claim=(
                "Trajectory analysis recovered branching blood lineages rooted "
                "in HSCs."
            ),
            confidence=0.82,
            claim_type="trajectory",
            evidence_steps=_evidence_steps(obs, {OutputType.TRAJECTORY_RESULT}),
        ))
    if "GATA1" in network_regulators:
        claims.append(ConclusionClaim(
            claim="GATA1 emerges as a driver of erythroid fate commitment.",
            confidence=0.8,
            claim_type="regulatory_network",
            evidence_steps=_evidence_steps(obs, {OutputType.NETWORK_RESULT}),
        ))
    if {"CEBPA", "SPI1"} & network_regulators:
        claims.append(ConclusionClaim(
            claim="CEBPA and SPI1 support myeloid branch decisions.",
            confidence=0.78,
            claim_type="regulatory_network",
            evidence_steps=_evidence_steps(obs, {OutputType.NETWORK_RESULT}),
        ))

    return claims


def compare_expected_findings(
    task: TaskSpec,
    obs: ExperimentObservation,
) -> tuple[List[str], List[str]]:
    """Compare the episode evidence against literature-backed findings."""

    evidence_text = _evidence_text(obs)
    matched: List[str] = []
    missed: List[str] = []

    for finding in task.expected_findings:
        keywords = [kw.lower() for kw in finding.keywords]
        if not keywords:
            keywords = _tokenize(finding.finding)
        hits = sum(1 for kw in keywords if kw in evidence_text)
        threshold = max(1, (len(keywords) + 1) // 2)
        if hits >= threshold:
            matched.append(finding.finding)
        else:
            missed.append(finding.finding)

    return matched, missed


def _resolve_scenario(
    problem_statement: str,
    scenario_name: Optional[str],
) -> Scenario:
    if scenario_name:
        for scenario in SCENARIO_LIBRARY:
            if scenario.name == scenario_name:
                return scenario
        raise ValueError(f"Unknown scenario_name '{scenario_name}'.")
    return select_literature_scenario(problem_statement)


def _tool_context(
    task: TaskSpec,
    *,
    libraries: Sequence[str],
    include_expected_findings: bool = False,
) -> Dict[str, Any]:
    context: Dict[str, Any] = {
        "literature_query": task.problem_statement,
        "paper_references": [
            {
                "title": ref.title,
                "doi": ref.doi,
                "pmid": ref.pmid,
                "url": ref.url,
            }
            for ref in task.paper_references
        ],
        "bioinformatics_libraries": list(libraries),
    }
    if include_expected_findings:
        context["expected_findings"] = [
            finding.finding for finding in task.expected_findings
        ]
    return context


def _default_comparison_name(task: TaskSpec) -> str:
    conditions = {condition.lower() for condition in task.conditions}
    if {"healthy", "ipf"} <= conditions:
        return "IPF_vs_healthy"
    if any("treated" in condition for condition in conditions) and any(
        "untreated" in condition for condition in conditions
    ):
        return "treated_vs_untreated"
    if any("healthy" in condition for condition in conditions):
        return "disease_vs_healthy"
    return "disease_vs_healthy"


def _preferred_marker(task: TaskSpec) -> str:
    """Derive a candidate marker from the problem statement, not expected findings."""
    tokens = [t for t in TOKEN_RE.findall(task.problem_statement) if t.isupper() and len(t) >= 3]
    if tokens:
        return tokens[0]
    return "unknown"


def _latest_output_data(
    obs: ExperimentObservation,
    output_type: OutputType,
) -> Dict[str, Any]:
    for output in reversed(obs.all_outputs):
        if output.output_type == output_type:
            return output.data
    return {}


def _extract_network_regulators(obs: ExperimentObservation) -> List[str]:
    for output in reversed(obs.all_outputs):
        if output.output_type == OutputType.NETWORK_RESULT:
            return output.data.get("top_regulators", [])
    return []


def _evidence_steps(
    obs: ExperimentObservation,
    output_types: set[OutputType],
) -> List[int]:
    return [
        output.step_index
        for output in obs.all_outputs
        if output.output_type in output_types
    ]


def _evidence_text(obs: ExperimentObservation) -> str:
    parts: List[str] = []
    parts.extend(obs.discovered_markers)
    parts.extend(obs.candidate_mechanisms)
    parts.extend(conclusion.claim for conclusion in obs.conclusions)

    for output in obs.all_outputs:
        parts.append(output.summary)
        if output.output_type == OutputType.DE_RESULT:
            parts.extend(
                gene["gene"]
                for gene in output.data.get("top_genes", [])
                if isinstance(gene, dict) and "gene" in gene
            )
        elif output.output_type == OutputType.PATHWAY_RESULT:
            parts.extend(
                pathway["pathway"]
                for pathway in output.data.get("top_pathways", [])
                if isinstance(pathway, dict) and "pathway" in pathway
            )
        elif output.output_type == OutputType.NETWORK_RESULT:
            parts.extend(output.data.get("top_regulators", []))
        elif output.output_type == OutputType.TRAJECTORY_RESULT:
            if output.data.get("branching_detected"):
                parts.append("branching lineage HSC trajectory")

    return " ".join(parts).lower()


def _tokenize(text: str) -> List[str]:
    return [
        token.lower()
        for token in TOKEN_RE.findall(text)
        if token and token.lower() not in STOPWORDS
    ]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--problem-statement",
        default=(
            "Design a follow-up validation experiment for candidate biomarker "
            "SPP1 in idiopathic pulmonary fibrosis."
        ),
    )
    parser.add_argument("--scenario-name", default=None)
    parser.add_argument("--domain-randomise", action="store_true")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    result = run_paper_benchmark(
        problem_statement=args.problem_statement,
        scenario_name=args.scenario_name,
        domain_randomise=args.domain_randomise,
    )

    if args.json:
        print(json.dumps(result.to_dict(), indent=2))
        return

    print(f"Scenario: {result.scenario_name}")
    print(f"Problem: {result.problem_statement}")
    print(f"Paper: {', '.join(result.matched_papers)}")
    print(f"Match ratio: {result.match_ratio:.2%}")
    print(f"Matched findings: {len(result.matched_findings)}")
    print(f"Missed findings: {len(result.missed_findings)}")
    print(f"Discovered markers: {', '.join(result.discovered_markers[:8])}")
    print(f"Candidate mechanisms: {', '.join(result.candidate_mechanisms[:5])}")
    print(f"Conclusions: {len(result.conclusions)}")
    print(f"Final reward: {result.final_reward:+.3f}")
    print(f"Bio libraries: {json.dumps(result.bio_library_versions, sort_keys=True)}")


if __name__ == "__main__":
    main()