src/api/routes.py

"""POST /pipeline/{bbb,eeg,mri} routes — thin dispatchers over the pipelines.

Each route validates its request body via Pydantic, invokes the pipeline,
reads back the produced Parquet to populate row/column counts, and returns
a uniform PipelineResponse. Pipeline-domain errors map to standard HTTP
codes: FileNotFoundError -> 404, ValueError -> 400, anything else -> 500.
"""
from __future__ import annotations

import os
import time
from collections import deque
from pathlib import Path
from typing import Callable

import mlflow
import pandas as pd
from fastapi import APIRouter, HTTPException

from src.api.schemas import (
    BBBExplainRequest,
    BBBExplainResponse,
    BBBPredictRequest,
    BBBPredictResponse,
    BBBRequest,
    CalibrationContext,
    EEGExplainRequest,
    EEGExplainResponse,
    EEGRequest,
    FeatureAttribution,
    HarmonizationRow,
    MLflowRunsResponse,
    MLflowRunSummary,
    ModelProvenance,
    MRIDiagnosticsRequest,
    MRIDiagnosticsResponse,
    MRIExplainRequest,
    MRIExplainResponse,
    MRIRequest,
    PipelineResponse,
    RunDiffRequest,
    RunDiffResponse,
    RunDiffRow,
)
from src.core.logger import get_logger
from src.llm import explainer as llm_explainer
from src.models import bbb_model
from src.pipelines import bbb_pipeline, eeg_pipeline, mri_pipeline

logger = get_logger(__name__)
router = APIRouter(prefix="/pipeline")
predict_router = APIRouter(prefix="/predict")
explain_router = APIRouter(prefix="/explain")
experiments_router = APIRouter(prefix="/experiments")


def _wrap(
    experiment_name: str,
    output_path: Path,
    fn: Callable[[], None],
) -> PipelineResponse:
    """Run `fn()` (the pipeline call), gather metrics, return PipelineResponse."""
    started = time.perf_counter()
    try:
        fn()
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except (ValueError, KeyError) as e:
        # KeyError: MRI pipeline raises this when sites_csv is missing a
        # site assignment for a subject — a user data problem, not a 500.
        raise HTTPException(status_code=400, detail=str(e))
    duration_sec = time.perf_counter() - started

    df = pd.read_parquet(output_path)
    runs = mlflow.search_runs(
        experiment_names=[experiment_name],
        max_results=1,
        order_by=["start_time DESC"],
    )
    run_id = runs.iloc[0]["run_id"] if len(runs) else None

    return PipelineResponse(
        status="ok",
        output_path=str(output_path),
        rows=len(df),
        columns=df.shape[1],
        duration_sec=duration_sec,
        mlflow_run_id=run_id,
    )


@router.post("/bbb", response_model=PipelineResponse)
def run_bbb(req: BBBRequest) -> PipelineResponse:
    """Run the BBB pipeline; return rows/cols/duration + the MLflow run id."""
    return _wrap(
        "bbb_pipeline",
        Path(req.output_path),
        lambda: bbb_pipeline.run_pipeline(
            input_path=Path(req.input_path),
            output_path=Path(req.output_path),
            smiles_col=req.smiles_col,
            n_bits=req.n_bits,
            radius=req.radius,
        ),
    )


@router.post("/eeg", response_model=PipelineResponse)
def run_eeg(req: EEGRequest) -> PipelineResponse:
    """Run the EEG pipeline; return rows/cols/duration + the MLflow run id."""
    return _wrap(
        "eeg_pipeline",
        Path(req.output_path),
        lambda: eeg_pipeline.run_pipeline(
            input_path=Path(req.input_path),
            output_path=Path(req.output_path),
            epoch_duration_s=req.epoch_duration_s,
            eog_ch_name=req.eog_ch_name,
            n_components=req.n_components,
            random_state=req.random_state,
        ),
    )


@router.post("/mri", response_model=PipelineResponse)
def run_mri(req: MRIRequest) -> PipelineResponse:
    """Run the MRI pipeline; return rows/cols/duration + the MLflow run id."""
    return _wrap(
        "mri_pipeline",
        Path(req.output_path),
        lambda: mri_pipeline.run_pipeline(
            input_dir=Path(req.input_dir),
            sites_csv=Path(req.sites_csv),
            output_path=Path(req.output_path),
        ),
    )


# Default artifact location. Overridable via BBB_MODEL_PATH env var so tests
# can point at a tmp-built model without touching production paths.
_DEFAULT_BBB_MODEL_PATH = Path("data/processed/bbb_model.joblib")


def _bbb_model_path() -> Path:
    """Return the BBB model artifact path, overridable via BBB_MODEL_PATH env var."""
    return Path(os.environ.get("BBB_MODEL_PATH", str(_DEFAULT_BBB_MODEL_PATH)))


# Per-worker rolling window of recent prediction confidences.
# Cleared on worker restart; multi-worker setups have independent windows.
WORKER_CONFIDENCE_DEQUE: deque[float] = deque(maxlen=100)
_DRIFT_MIN_SAMPLES = 10


def _compute_drift_z(model, confidence: float) -> tuple[float | None, int]:
    """Append `confidence` to the worker deque and compute the drift z-score.

    Returns (drift_z, rolling_n). drift_z is None until both:
      (1) the deque has at least `_DRIFT_MIN_SAMPLES` samples, AND
      (2) the model has `_neurobridge_train_stats` attached.

    z = (rolling_median - train_median) / max(train_std, 1e-9)
    """
    import statistics

    WORKER_CONFIDENCE_DEQUE.append(float(confidence))
    rolling_n = len(WORKER_CONFIDENCE_DEQUE)
    stats = getattr(model, "_neurobridge_train_stats", None)
    if rolling_n < _DRIFT_MIN_SAMPLES or stats is None:
        return None, rolling_n
    rolling_median = statistics.median(WORKER_CONFIDENCE_DEQUE)
    train_median = float(stats["median"])
    train_std = max(float(stats["std"]), 1e-9)
    drift_z = (rolling_median - train_median) / train_std
    return float(drift_z), rolling_n


_PROVENANCE_CACHE: ModelProvenance | None = None
_MODEL_VERSION = "v1"  # bump manually per train cycle


def _build_provenance(model) -> ModelProvenance:
    """Look up the most recent BBB MLflow run; build a ModelProvenance.

    Cached at module level so we hit MLflow once per worker. Failures (no
    runs found, MLflow unreachable, NEUROBRIDGE_DISABLE_MLFLOW=1) all
    degrade to a partial ModelProvenance with mlflow_run_id=None — the
    badge still renders, just without a run id.
    """
    global _PROVENANCE_CACHE
    if _PROVENANCE_CACHE is not None:
        # Refresh n_examples each call from the model (cheap lookup).
        n_train = None
        stats = getattr(model, "_neurobridge_train_stats", None)
        if stats is not None:
            n_train = int(stats.get("n_train", 0)) or None
        return _PROVENANCE_CACHE.model_copy(update={"n_examples": n_train})

    run_id: str | None = None
    train_date: str | None = None
    if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") != "1":
        try:
            runs = mlflow.search_runs(
                experiment_names=["bbb_pipeline"],
                max_results=1,
                order_by=["start_time DESC"],
            )
            if len(runs):
                row = runs.iloc[0]
                run_id = str(row["run_id"])
                ts = row.get("start_time")
                if ts is not None:
                    train_date = str(pd.Timestamp(ts).isoformat())
        except Exception as e:  # broad: MLflow store unreachable, schema mismatch, etc.
            logger.warning("MLflow provenance lookup failed: %s", e)

    n_train = None
    stats = getattr(model, "_neurobridge_train_stats", None)
    if stats is not None:
        n_train = int(stats.get("n_train", 0)) or None

    _PROVENANCE_CACHE = ModelProvenance(
        mlflow_run_id=run_id,
        model_version=_MODEL_VERSION,
        train_date=train_date,
        n_examples=n_train,
    )
    return _PROVENANCE_CACHE


def _matching_calibration_bin(model, confidence: float) -> CalibrationContext | None:
    """Pick the highest-threshold bin whose threshold <= confidence. None if no match or no metadata."""
    bins = getattr(model, "_neurobridge_calibration", None)
    if not bins:
        return None
    matched = None
    for bin_ in bins:
        if bin_["threshold"] <= confidence:
            matched = bin_
        else:
            break
    if matched is None:
        return None
    return CalibrationContext(
        threshold=matched["threshold"],
        precision=matched["precision"],
        support=matched["support"],
    )


@predict_router.post("/bbb", response_model=BBBPredictResponse)
def predict_bbb(req: BBBPredictRequest) -> BBBPredictResponse:
    """Predict BBB permeability + return SHAP attributions for one SMILES.

    Returns 503 if the model artifact is missing (operator hasn't run the
    trainer CLI yet); 400 on invalid SMILES; 200 with the decision payload
    on success.
    """
    artifact = _bbb_model_path()
    if not artifact.exists():
        raise HTTPException(
            status_code=503,
            detail=(
                f"BBB model artifact not available at {artifact}. "
                f"Run `python -m src.models.bbb_model` to train it."
            ),
        )
    try:
        model = bbb_model.load(artifact)
    except FileNotFoundError as e:
        raise HTTPException(status_code=503, detail=str(e))

    try:
        pred = bbb_model.predict_with_proba(model, req.smiles)
        attributions = bbb_model.explain_prediction(model, req.smiles, top_k=req.top_k)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    label_text = "permeable" if pred["label"] == 1 else "non-permeable"
    calibration = _matching_calibration_bin(model, pred["confidence"])
    drift_z, rolling_n = _compute_drift_z(model, pred["confidence"])
    provenance = _build_provenance(model)
    return BBBPredictResponse(
        label=pred["label"],
        label_text=label_text,
        confidence=pred["confidence"],
        top_features=[FeatureAttribution(**a) for a in attributions],
        calibration=calibration,
        drift_z=drift_z,
        rolling_n=rolling_n,
        provenance=provenance,
    )


@router.post("/mri/diagnostics", response_model=MRIDiagnosticsResponse)
def mri_diagnostics(req: MRIDiagnosticsRequest) -> MRIDiagnosticsResponse:
    """Run the MRI pipeline twice and return pre/post ComBat data + site-gap KPIs."""
    input_dir = Path(req.input_dir)
    sites_csv = Path(req.sites_csv)
    try:
        df = mri_pipeline.compute_harmonization_diagnostics(
            input_dir=input_dir, sites_csv=sites_csv,
        )
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except KeyError as e:
        raise HTTPException(status_code=400, detail=str(e))

    if df.empty:
        return MRIDiagnosticsResponse(
            rows=[], site_gap_pre=0.0, site_gap_post=0.0, reduction_factor=0.0,
        )

    # Site-gap KPI on the first feature, averaged per site
    feat = df["feature"].iloc[0]
    feat_df = df[df["feature"] == feat]
    pre_means = feat_df[feat_df["harmonization_state"] == "Pre-ComBat"].groupby(
        "site"
    )["feature_value"].mean()
    post_means = feat_df[feat_df["harmonization_state"] == "Post-ComBat"].groupby(
        "site"
    )["feature_value"].mean()
    site_gap_pre = float(pre_means.max() - pre_means.min())
    site_gap_post = float(post_means.max() - post_means.min())
    eps = 1e-9
    reduction_factor = site_gap_pre / max(site_gap_post, eps)

    rows = [
        HarmonizationRow(**rec) for rec in df.to_dict(orient="records")
    ]
    return MRIDiagnosticsResponse(
        rows=rows,
        site_gap_pre=site_gap_pre,
        site_gap_post=site_gap_post,
        reduction_factor=reduction_factor,
    )


@explain_router.post("/bbb", response_model=BBBExplainResponse)
def explain_bbb(req: BBBExplainRequest) -> BBBExplainResponse:
    """Natural-language rationale for a single BBB prediction.

    Always returns 200 — the explainer is guaranteed to produce a
    rationale via deterministic-template fallback. Pydantic enforces
    a non-empty top_features list; an empty list returns 422 from
    FastAPI before this handler runs.
    """
    payload: llm_explainer.ExplainPayload = {
        "smiles": req.smiles,
        "label": req.label,
        "label_text": req.label_text,
        "confidence": req.confidence,
        "top_features": [
            {"feature": f.feature, "shap_value": f.shap_value}
            for f in req.top_features
        ],
        "calibration": (
            None
            if req.calibration is None
            else {
                "threshold": req.calibration.threshold,
                "precision": req.calibration.precision,
                "support": req.calibration.support,
            }
        ),
        "drift_z": req.drift_z,
        "user_question": req.user_question or "",
    }
    result = llm_explainer.explain(payload)
    return BBBExplainResponse(
        rationale=result["rationale"],
        source=result["source"],
        model=result["model"],
    )


@explain_router.post("/eeg", response_model=EEGExplainResponse)
def explain_eeg(req: EEGExplainRequest) -> EEGExplainResponse:
    """Natural-language rationale for an EEG pipeline run."""
    payload = {
        "rows": req.rows,
        "columns": req.columns,
        "duration_sec": req.duration_sec,
        "mlflow_run_id": req.mlflow_run_id,
        "user_question": req.user_question or "",
    }
    result = llm_explainer.explain(payload, modality="eeg")
    return EEGExplainResponse(
        rationale=result["rationale"],
        source=result["source"],
        model=result["model"],
    )


@explain_router.post("/mri", response_model=MRIExplainResponse)
def explain_mri(req: MRIExplainRequest) -> MRIExplainResponse:
    """Natural-language rationale for an MRI ComBat diagnostic run."""
    payload = {
        "site_gap_pre": req.site_gap_pre,
        "site_gap_post": req.site_gap_post,
        "reduction_factor": req.reduction_factor,
        "n_subjects": req.n_subjects,
        "user_question": req.user_question or "",
    }
    result = llm_explainer.explain(payload, modality="mri")
    return MRIExplainResponse(
        rationale=result["rationale"],
        source=result["source"],
        model=result["model"],
    )


@experiments_router.get("/runs", response_model=MLflowRunsResponse)
def list_runs(limit: int = 50) -> MLflowRunsResponse:
    """List recent MLflow runs across known experiments.

    Returns an empty list when MLflow is disabled or unreachable.
    """
    if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") == "1":
        return MLflowRunsResponse(runs=[])

    summaries: list[MLflowRunSummary] = []
    for exp_name in ("bbb_pipeline", "eeg_pipeline", "mri_pipeline"):
        try:
            df = mlflow.search_runs(
                experiment_names=[exp_name],
                max_results=limit,
                order_by=["start_time DESC"],
            )
        except Exception as e:  # broad: MLflow store unreachable / not found
            logger.warning("MLflow lookup failed for %s: %s", exp_name, e)
            continue
        for _, row in df.iterrows():
            metrics = {
                col[len("metrics."):]: float(row[col])
                for col in df.columns
                if col.startswith("metrics.") and pd.notna(row[col])
            }
            params = {
                col[len("params."):]: str(row[col])
                for col in df.columns
                if col.startswith("params.") and pd.notna(row[col])
            }
            summaries.append(
                MLflowRunSummary(
                    run_id=str(row["run_id"]),
                    experiment_name=exp_name,
                    start_time=str(pd.Timestamp(row["start_time"]).isoformat())
                    if pd.notna(row.get("start_time"))
                    else "",
                    status=str(row.get("status", "UNKNOWN")),
                    metrics=metrics,
                    params=params,
                )
            )
    summaries.sort(key=lambda s: s.start_time, reverse=True)
    return MLflowRunsResponse(runs=summaries[:limit])


@experiments_router.post("/diff", response_model=RunDiffResponse)
def diff_runs(req: RunDiffRequest) -> RunDiffResponse:
    """Side-by-side diff of two MLflow runs (metrics + params).

    Returns 404 if either run id is not found in the local MLflow store.
    Returns 200 with an empty rows list when MLflow is disabled.
    """
    if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") == "1":
        return RunDiffResponse(rows=[])

    try:
        run_a = mlflow.get_run(req.run_id_a)
        run_b = mlflow.get_run(req.run_id_b)
    except Exception as e:
        raise HTTPException(status_code=404, detail=f"Run not found: {e}")

    metrics_a = run_a.data.metrics
    metrics_b = run_b.data.metrics
    params_a = run_a.data.params
    params_b = run_b.data.params

    rows: list[RunDiffRow] = []
    for key in sorted(set(metrics_a) | set(metrics_b)):
        va = metrics_a.get(key)
        vb = metrics_b.get(key)
        rows.append(
            RunDiffRow(
                key=key, kind="metric",
                value_a=None if va is None else f"{va:.6g}",
                value_b=None if vb is None else f"{vb:.6g}",
                differs=(va != vb),
            )
        )
    for key in sorted(set(params_a) | set(params_b)):
        va = params_a.get(key)
        vb = params_b.get(key)
        rows.append(
            RunDiffRow(
                key=key, kind="param",
                value_a=va, value_b=vb, differs=(va != vb),
            )
        )
    return RunDiffResponse(rows=rows)