api/main.py

"""Versioned REST API for the TP53 discovery workflow."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any
from uuid import uuid4

import pandas as pd
from fastapi import BackgroundTasks, FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from utils.calculations import (
    DEFAULT_WEIGHTS,
    SCORING_VERSION,
    apply_weights,
    bmut_from_score,
    consensus_score,
    empirical_docking_score,
    ingest_variant_frame,
    md_metrics,
    pocket_hydrophobicity,
    pocket_records,
    selectivity_from_scores,
    structure_record,
)
from api import store

DATA_DIR = Path(__file__).resolve().parent.parent / "data"

app = FastAPI(
    title="TP53 Mutant Discovery API",
    version="1.0.0",
    description="Traceable variant-to-compound analysis with asynchronous jobs and versioned artifacts.",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)


class ProjectCreate(BaseModel):
    name: str = Field(min_length=1)
    cancer_type: str = "Lung Cancer"
    metadata: dict[str, Any] = Field(default_factory=dict)


class VariantIngest(BaseModel):
    project_id: str
    records: list[dict[str, Any]]


class VariantAnalyze(BaseModel):
    project_id: str


class StructurePrepare(BaseModel):
    project_id: str
    variant_ids: list[str]
    source_preference: str = "experimental_then_template"


class PocketDetect(BaseModel):
    project_id: str
    structure_ids: list[str]


class CompoundScreen(BaseModel):
    project_id: str
    variant_id: str
    library_id: str = "platform-library-v1"
    model_id: str = "complementarity-v1"
    limit: int = Field(default=500, ge=1, le=2000)


class DockingRun(BaseModel):
    project_id: str
    variant_id: str
    compounds: list[dict[str, Any]]
    protocol_id: str = "matched-wt-mutant-v1"


class MDRun(BaseModel):
    project_id: str
    complexes: list[dict[str, Any]]
    protocol_id: str = "explicit-solvent-3x100ns-v1"


class RankingCompute(BaseModel):
    project_id: str
    variant_id: str
    candidates: list[dict[str, Any]]
    weights: dict[str, float] = Field(default_factory=lambda: DEFAULT_WEIGHTS.copy())


@app.on_event("startup")
def startup() -> None:
    store.init_db()


def require_project(project_id: str) -> dict:
    project = store.get_project(project_id)
    if not project:
        raise HTTPException(404, f"Project {project_id} not found")
    return project


def submit(background: BackgroundTasks, project_id: str, operation: str, payload: dict) -> dict:
    require_project(project_id)
    job = store.create_job(project_id, operation, payload)
    background.add_task(process_job, job["job_id"])
    return {"job_id": job["job_id"], "status": "queued", "status_url": f"/jobs/{job['job_id']}"}


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "version": app.version, "scoring_version": SCORING_VERSION}


@app.post("/projects", status_code=201)
def create_project(body: ProjectCreate) -> dict:
    return store.create_project(body.name, body.cancer_type, body.metadata)


@app.get("/projects/{project_id}")
def get_project(project_id: str) -> dict:
    return require_project(project_id)


@app.post("/variants:ingest", status_code=202)
def variants_ingest(body: VariantIngest, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "variants:ingest", body.model_dump())


@app.post("/variants/{variant_id}/analyze", status_code=202)
def variant_analyze(variant_id: str, body: VariantAnalyze, background: BackgroundTasks) -> dict:
    return submit(
        background,
        body.project_id,
        "variant:analyze",
        {**body.model_dump(), "variant_id": variant_id},
    )


@app.post("/structures:prepare", status_code=202)
def structures_prepare(body: StructurePrepare, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "structures:prepare", body.model_dump())


@app.post("/pockets:detect", status_code=202)
def pockets_detect(body: PocketDetect, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "pockets:detect", body.model_dump())


@app.post("/compounds:screen", status_code=202)
def compounds_screen(body: CompoundScreen, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "compounds:screen", body.model_dump())


@app.post("/docking:run", status_code=202)
def docking_run(body: DockingRun, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "docking:run", body.model_dump())


@app.post("/md:run", status_code=202)
def md_run(body: MDRun, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "md:run", body.model_dump())


@app.post("/rankings:compute", status_code=202)
def rankings_compute(body: RankingCompute, background: BackgroundTasks) -> dict:
    return submit(background, body.project_id, "rankings:compute", body.model_dump())


@app.get("/jobs/{job_id}")
def get_job(job_id: str) -> dict:
    job = store.get_job(job_id)
    if not job:
        raise HTTPException(404, f"Job {job_id} not found")
    return job


@app.get("/projects/{project_id}/jobs")
def project_jobs(project_id: str) -> list[dict]:
    require_project(project_id)
    return store.list_jobs(project_id)


@app.get("/projects/{project_id}/artifacts")
def project_artifacts(project_id: str) -> list[dict]:
    require_project(project_id)
    return store.list_artifacts(project_id)


@app.get("/artifacts/{artifact_id}")
def artifact(artifact_id: str) -> dict:
    row = store.get_artifact(artifact_id)
    if not row:
        raise HTTPException(404, f"Artifact {artifact_id} not found")
    return row


@app.get("/projects/{project_id}/report")
def project_report(project_id: str) -> dict:
    project = require_project(project_id)
    variants = store.list_records(project_id, "variant")
    structures = store.list_records(project_id, "structure")
    pockets = store.list_records(project_id, "pocket")
    rankings = store.list_records(project_id, "ranking")
    return {
        "project": project,
        "summary": {
            "variants": len(variants),
            "structures": len(structures),
            "pockets": len(pockets),
            "ranked_candidates": len(rankings),
        },
        "top_candidates": sorted(rankings, key=lambda x: x.get("rescue_score", 0), reverse=True)[:10],
        "jobs": store.list_jobs(project_id),
        "artifacts": store.list_artifacts(project_id),
        "versions": {"api": app.version, "scoring": SCORING_VERSION},
    }


def process_job(job_id: str) -> None:
    job = store.get_job(job_id)
    if not job:
        return
    try:
        store.update_job(job_id, "running", log=f"Started {job['operation']}")
        handler = HANDLERS[job["operation"]]
        result = handler(job["project_id"], job["input"], job_id)
        store.update_job(job_id, "succeeded", output=result, log="Completed successfully")
    except Exception as exc:
        store.update_job(job_id, "failed", error=str(exc), log=f"Failed: {exc}")


def handle_ingest(project_id: str, payload: dict, job_id: str) -> dict:
    frame = pd.DataFrame(payload["records"])
    normalized = ingest_variant_frame(frame, require_project(project_id)["cancer_type"])
    records = []
    for _, row in normalized.iterrows():
        item = _jsonable(row.to_dict())
        item["variant_id"] = f"VAR-{uuid4().hex[:12]}"
        records.append(item)
    ids = store.add_records(project_id, "variant", records)
    artifact = store.add_artifact(
        project_id, "qc-report", "variant-qc.json", records, "application/json", job_id
    )
    return {"variant_ids": ids, "normalized_records": records, "qc_artifact_id": artifact["artifact_id"]}


def handle_analyze(project_id: str, payload: dict, job_id: str) -> dict:
    record = store.get_record(project_id, "variant", payload["variant_id"])
    if not record:
        raise ValueError(f"Variant {payload['variant_id']} not found")
    return {
        "variant_id": payload["variant_id"],
        "hgvs_p": record["hgvs_p"],
        "priority": record["priority_score"],
        "class": record["functional_class"],
        "domain": record["domain"],
        "route": record["route"],
        "confidence": record["confidence"],
        "qc_flags": record["qc_flags"],
    }


def handle_structures(project_id: str, payload: dict, job_id: str) -> dict:
    rows = []
    for variant_id in payload["variant_ids"]:
        variant = store.get_record(project_id, "variant", variant_id)
        if not variant:
            raise ValueError(f"Variant {variant_id} not found")
        rec = _jsonable(structure_record(pd.Series(variant)))
        rec["structure_id"] = f"STR-{uuid4().hex[:12]}"
        rec["variant_id"] = variant_id
        rec["source_preference"] = payload["source_preference"]
        rows.append(rec)
    ids = store.add_records(project_id, "structure", rows)
    artifact = store.add_artifact(project_id, "structures", "structures.json", rows, "application/json", job_id)
    return {"structure_ids": ids, "structures": rows, "artifact_id": artifact["artifact_id"]}


def handle_pockets(project_id: str, payload: dict, job_id: str) -> dict:
    rows = []
    for structure_id in payload["structure_ids"]:
        structure = store.get_record(project_id, "structure", structure_id)
        if not structure:
            raise ValueError(f"Structure {structure_id} not found")
        for rec in pocket_records(structure):
            rec["pocket_id"] = f"POC-{uuid4().hex[:12]}"
            rec["structure_id"] = structure_id
            rows.append(_jsonable(rec))
    ids = store.add_records(project_id, "pocket", rows)
    artifact = store.add_artifact(project_id, "pockets", "pockets.json", rows, "application/json", job_id)
    return {"pocket_ids": ids, "pockets": rows, "artifact_id": artifact["artifact_id"]}


def handle_screen(project_id: str, payload: dict, job_id: str) -> dict:
    variant = store.get_record(project_id, "variant", payload["variant_id"])
    if not variant:
        raise ValueError(f"Variant {payload['variant_id']} not found")
    source = pd.read_csv(DATA_DIR / "screening.csv")
    rows = source[source["variant"] == variant["hgvs_p"]].nlargest(payload["limit"], "ai_score")
    output = [_jsonable(row) for row in rows.to_dict("records")]
    artifact = store.add_artifact(project_id, "screening", "screening.json", output, "application/json", job_id)
    return {
        "variant_id": payload["variant_id"],
        "model_id": payload["model_id"],
        "library_id": payload["library_id"],
        "candidates": output,
        "artifact_id": artifact["artifact_id"],
    }


def handle_docking(project_id: str, payload: dict, job_id: str) -> dict:
    variant = store.get_record(project_id, "variant", payload["variant_id"])
    if not variant:
        raise ValueError(f"Variant {payload['variant_id']} not found")
    structure = next(
        (s for s in store.list_records(project_id, "structure") if s["variant_id"] == payload["variant_id"]),
        None,
    )
    if not structure:
        structure = _jsonable(structure_record(pd.Series(variant)))
    hyd = pocket_hydrophobicity(int(variant["position"]), variant["hgvs_p"])
    rows = []
    for compound in payload["compounds"]:
        required = ("compound_id", "mw", "logp", "tpsa", "rotbonds")
        missing = [key for key in required if key not in compound]
        if missing:
            raise ValueError(f"Compound missing {missing}")
        dmut = empirical_docking_score(
            compound["mw"], compound["logp"], compound["tpsa"], compound["rotbonds"],
            structure.get("pocket_vol_mut") or 165, hyd,
        )
        dwt = empirical_docking_score(
            compound["mw"], compound["logp"], compound["tpsa"], compound["rotbonds"],
            structure.get("pocket_vol_wt") or 165, 0.30, wt_pocket=True,
        )
        rows.append({
            **compound,
            "docking_id": f"DOCK-{uuid4().hex[:12]}",
            "variant_id": payload["variant_id"],
            "dock_mut": round(dmut, 3),
            "dock_wt": round(dwt, 3),
            "delta_dock": round(dmut - dwt, 3),
            "mutant_preference": round(selectivity_from_scores(dmut, dwt), 3),
            "protocol_version": payload["protocol_id"],
            "units": "kcal/mol",
        })
    ids = store.add_records(project_id, "docking", rows)
    artifact = store.add_artifact(project_id, "docking", "docking.json", rows, "application/json", job_id)
    return {"docking_ids": ids, "poses": rows, "artifact_id": artifact["artifact_id"]}


def handle_md(project_id: str, payload: dict, job_id: str) -> dict:
    rows = []
    for complex_data in payload["complexes"]:
        for key in ("complex_id", "dock_mut"):
            if key not in complex_data:
                raise ValueError(f"Complex missing {key}")
        bmut = bmut_from_score(complex_data["dock_mut"])
        selectivity = float(complex_data.get("mutant_preference", 0))
        metrics = md_metrics(
            complex_data["dock_mut"], complex_data.get("rotbonds", 4), bmut, selectivity
        )
        rows.append({
            "md_id": f"MD-{uuid4().hex[:12]}",
            **complex_data,
            **metrics,
            "replicas": 3,
            "duration_ns": 100,
            "protocol_version": payload["protocol_id"],
        })
    ids = store.add_records(project_id, "md", rows)
    artifact = store.add_artifact(project_id, "md", "md-summary.json", rows, "application/json", job_id)
    return {"md_ids": ids, "results": rows, "artifact_id": artifact["artifact_id"]}


def handle_rankings(project_id: str, payload: dict, job_id: str) -> dict:
    candidates = pd.DataFrame(payload["candidates"])
    required = list(DEFAULT_WEIGHTS)
    missing = [col for col in required if col not in candidates.columns]
    if missing:
        raise ValueError(f"Candidates missing score components: {missing}")
    if "variant" not in candidates:
        candidates["variant"] = payload["variant_id"]
    if "confidence" not in candidates:
        candidates["confidence"] = 0.5
    if "admet_flag" not in candidates:
        candidates["admet_flag"] = "Watch"
    ranked = apply_weights(candidates, payload["weights"])
    rows = []
    for rec in ranked.to_dict("records"):
        rec = _jsonable(rec)
        rec["ranking_id"] = f"RNK-{uuid4().hex[:12]}"
        rec["variant_id"] = payload["variant_id"]
        rec["scoring_version"] = SCORING_VERSION
        rows.append(rec)
    ids = store.add_records(project_id, "ranking", rows)
    artifact = store.add_artifact(project_id, "ranking", "ranking.json", rows, "application/json", job_id)
    return {"ranking_ids": ids, "shortlist": rows[:10], "artifact_id": artifact["artifact_id"]}


HANDLERS = {
    "variants:ingest": handle_ingest,
    "variant:analyze": handle_analyze,
    "structures:prepare": handle_structures,
    "pockets:detect": handle_pockets,
    "compounds:screen": handle_screen,
    "docking:run": handle_docking,
    "md:run": handle_md,
    "rankings:compute": handle_rankings,
}


def _jsonable(value: Any) -> Any:
    if isinstance(value, dict):
        return {k: _jsonable(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_jsonable(v) for v in value]
    if pd.isna(value):
        return None
    if hasattr(value, "item"):
        return value.item()
    return value


# Mounted last so API routes take precedence; Docker serves API and dashboard together.
app.mount("/", StaticFiles(directory=Path(__file__).resolve().parent.parent, html=True), name="dashboard")