api/main.py

"""NephroScreen FastAPI service.

Loads the train-fitted pipeline once at startup and exposes:
  GET  /health          liveness + model status
  GET  /api/metadata    field schema (labels, choices, ranges) for the frontend
  POST /api/predict     CKD probability, risk band, and rule-based indicators
The static frontend is served at / when the frontend/ directory is present.
"""
import json
from contextlib import asynccontextmanager
from pathlib import Path

import joblib
import pandas as pd
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

from nephroscreen.config import (
    CATEGORICAL_CHOICES,
    FEATURE_LABELS,
    METRICS_PATH,
    MODEL_PATH,
    NUMERIC_COLUMNS,
    RAW_INPUT_COLUMNS,
    REFERENCE_RANGES,
)

from .schemas import Indicator, PatientInput, PredictionResponse

DISCLAIMER = (
    "Educational demo only. Trained on a small public dataset (UCI CKD, 400 "
    "patients) and NOT validated for clinical use. Do not use for diagnosis."
)

_model = None
_manifest = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _model, _manifest
    if MODEL_PATH.exists():
        _model = joblib.load(MODEL_PATH)
    if METRICS_PATH.exists():
        _manifest = json.loads(METRICS_PATH.read_text(encoding="utf-8"))
    yield


app = FastAPI(
    title="NephroScreen API",
    version="1.0.0",
    description="Early chronic kidney disease screening from routine lab values.",
    lifespan=lifespan,
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _model is not None}


@app.get("/api/metadata")
def metadata():
    """Everything the frontend needs to render the form and show model context."""
    return {
        "numeric_fields": [
            {"name": c, "label": FEATURE_LABELS[c]} for c in NUMERIC_COLUMNS
        ],
        "categorical_fields": [
            {"name": c, "label": FEATURE_LABELS[c], "choices": CATEGORICAL_CHOICES[c]}
            for c in CATEGORICAL_CHOICES
        ],
        "threshold": _manifest.get("serving_threshold", 0.5),
        "metrics": _manifest.get("metrics", {}),
        "top_features": _manifest.get("top_features", {}),
        "note": _manifest.get("note", ""),
    }


def _risk_band(prob: float) -> str:
    if prob >= 0.66:
        return "High"
    if prob >= 0.33:
        return "Moderate"
    return "Low"


def _indicators(data: dict) -> list[Indicator]:
    out = []
    for feat, (low, high, unit) in REFERENCE_RANGES.items():
        if feat.endswith("_note"):
            continue
        val = data.get(feat)
        if val is None:
            continue
        flag = "low" if val < low else "high" if val > high else "normal"
        if flag == "normal":
            continue  # surface only the abnormal drivers
        out.append(
            Indicator(
                feature=feat,
                label=FEATURE_LABELS.get(feat, feat),
                value=float(val),
                normal_range=f"{low}-{high} {unit}",
                flag=flag,
            )
        )
    return out


@app.post("/api/predict", response_model=PredictionResponse)
def predict(patient: PatientInput):
    if _model is None:
        raise HTTPException(503, "Model not loaded. Run training first.")

    data = patient.model_dump()
    row = {col: data.get(col) for col in RAW_INPUT_COLUMNS}
    frame = pd.DataFrame([row], columns=RAW_INPUT_COLUMNS)

    prob = float(_model.predict_proba(frame)[:, 1][0])
    threshold = float(_manifest.get("serving_threshold", 0.5))
    prediction = "CKD" if prob >= threshold else "Not CKD"

    return PredictionResponse(
        prediction=prediction,
        probability=round(prob, 4),
        risk_band=_risk_band(prob),
        threshold=threshold,
        key_indicators=_indicators(data),
        disclaimer=DISCLAIMER,
    )


# Serve the static frontend last so /api/* and /health take precedence.
_frontend = Path(__file__).resolve().parents[1] / "frontend"
if _frontend.exists():
    app.mount("/", StaticFiles(directory=str(_frontend), html=True), name="frontend")