Running benchmarks across datasets

"""Batch benchmarks: run one analysis across many datasets, then grade each
result in a second task."""

import json
import os
import shutil
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from phylo_client import PhyloClient

# Five public TCGA PanCancer Atlas studies on cBioPortal.
STUDIES = [
    "luad_tcga_pan_can_atlas_2018",
    "coadread_tcga_pan_can_atlas_2018",
    "brca_tcga_pan_can_atlas_2018",
    "skcm_tcga_pan_can_atlas_2018",
    "stad_tcga_pan_can_atlas_2018",
]

ANALYSIS_PROMPT = """
Test whether TP53 mutation status is associated with overall survival in the
cBioPortal study {study}. Download the study's mutation and clinical data from
cBioPortal yourself.

- Call a patient TP53-mutant if they carry any nonsynonymous TP53 mutation.
- Compare overall survival with Kaplan-Meier curves and a log-rank test.
- Fit a Cox model for TP53 status adjusted for age at diagnosis and AJCC
  pathologic stage.

Save two files under /mnt/results/:

1. report.md: cohort, definitions, methods, code, results, interpretation, and
   limitations.
2. result.json, raw JSON with exactly these keys: study_id, n_patients,
   n_tp53_mutant, logrank_p, cox_hr, cox_hr_ci_low, cox_hr_ci_high, cox_p,
   conclusion.
""".strip()

GRADING_PROMPT = """
Grade the TP53 overall-survival analysis of {study} in the project files
{prefix}report.md and {prefix}result.json. Judge only those two files. Do not
redo the analysis or read any other file.

Score each criterion and explain why:
1. Cohort and TP53 definition (20): nonsynonymous mutations only, patients
   without survival data excluded, counts reported.
2. Survival analysis (25): Kaplan-Meier plus log-rank test, done correctly.
3. Adjusted model (25): Cox model with age and stage, hazard ratio, 95% CI and
   p-value reported.
4. Consistency (15): result.json matches the numbers in report.md.
5. Interpretation (15): conclusion follows from the numbers, limitations noted.

Save raw JSON (no code fence) at /mnt/results/score.json:
{{"study_id": "{study}", "total_score": 0,
  "criteria": [{{"criterion": "...", "points": 0, "max_points": 0, "reason": "..."}}],
  "summary": "..."}}
""".strip()


def run_analysis(phylo, project_id, study):
    task = phylo.create_task(
        project_id=project_id,
        title=f"TP53 survival: {study}",
        initial_messages=[
            {"role": "user", "content": ANALYSIS_PROMPT.format(study=study)}
        ],
        model="max",
        disable_fallbacks=True,
        auto_mode=True,
    )
    print(f"[{study}] analysis task {task['id']}")
    phylo.wait_for_task(task["id"], poll_interval_seconds=60)
    return phylo.download_results(task["id"], Path("benchmark-results") / study)


def grade_analysis(phylo, project_id, study, results_dir):
    # Upload the analysis output under study-specific names so every grading
    # task in the shared project reads only its own files.
    prefix = f"{study}__"
    for name in ("report.md", "result.json"):
        staged = results_dir / f"{prefix}{name}"
        shutil.copyfile(results_dir / name, staged)
        phylo.upload_file(project_id, staged)

    task = phylo.create_task(
        project_id=project_id,
        title=f"Grade TP53 survival: {study}",
        initial_messages=[
            {
                "role": "user",
                "content": GRADING_PROMPT.format(study=study, prefix=prefix),
            }
        ],
        model="max",
        disable_fallbacks=True,
        auto_mode=True,
    )
    print(f"[{study}] grading task {task['id']}")
    phylo.wait_for_task(task["id"], poll_interval_seconds=60)
    grade_dir = phylo.download_results(task["id"], results_dir / "grade")
    return json.loads((grade_dir / "score.json").read_text())


def run_study(project_id, study):
    phylo = PhyloClient()  # one HTTP session per thread
    results_dir = run_analysis(phylo, project_id, study)
    score = grade_analysis(phylo, project_id, study, results_dir)
    result = json.loads((results_dir / "result.json").read_text())
    return study, result, score


def main():
    phylo = PhyloClient()
    project_id = (
        os.getenv("PHYLO_PROJECT_ID")
        or phylo.create_project(title="API demo")["id"]
    )

    # At most 10 tasks run at once; each study runs one task at a time.
    with ThreadPoolExecutor(max_workers=len(STUDIES)) as pool:
        outcomes = list(pool.map(lambda s: run_study(project_id, s), STUDIES))

    print(f"\n{'study':34} {'HR (95% CI)':22} {'cox p':>8} {'score':>6}")
    for study, result, score in outcomes:
        ci = (
            f"{result['cox_hr']:.2f} "
            f"({result['cox_hr_ci_low']:.2f}-{result['cox_hr_ci_high']:.2f})"
        )
        print(f"{study:34} {ci:22} {result['cox_p']:>8.3g} {score['total_score']:>6}")


if __name__ == "__main__":
    main()