"""Multimodal heart-attack risk app — ensemble router.

One ``POST /predict`` endpoint accepts multipart/form-data carrying *optional*
Framingham tabular fields and an *optional* ECG image.  An internal router picks
the path:

    tabular only   -> Model A (Framingham CHD)
    ECG only       -> Model B (ResNet ECG)
    both           -> A + B, averaged
    neither        -> HTTP 400

The predictor modules are imported lazily so the app still boots (and serves the
frontend) before the models have been trained.

Run:  uvicorn app:app --reload
"""
from __future__ import annotations

from pathlib import Path
from typing import Optional

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from inference.fusion import band_for, combine
from inference.validation import validate_tabular

# Canonical Framingham feature order (mirrors train_framingham.FEATURES).
FEATURES = [
    "male", "age", "education", "currentSmoker", "cigsPerDay", "BPMeds",
    "prevalentStroke", "prevalentHyp", "diabetes", "totChol", "sysBP",
    "diaBP", "BMI", "heartRate", "glucose",
]

app = FastAPI(title="Multimodal Heart Attack Risk — Ensemble Router")


def _has_value(v: Optional[str]) -> bool:
    return v is not None and str(v).strip() != ""


@app.post("/predict")
async def predict(
    # ── Tabular branch (all optional → blanks are KNN-imputed) ───────────────
    male: Optional[str] = Form(None),
    age: Optional[str] = Form(None),
    education: Optional[str] = Form(None),
    currentSmoker: Optional[str] = Form(None),
    cigsPerDay: Optional[str] = Form(None),
    BPMeds: Optional[str] = Form(None),
    prevalentStroke: Optional[str] = Form(None),
    prevalentHyp: Optional[str] = Form(None),
    diabetes: Optional[str] = Form(None),
    totChol: Optional[str] = Form(None),
    sysBP: Optional[str] = Form(None),
    diaBP: Optional[str] = Form(None),
    BMI: Optional[str] = Form(None),
    heartRate: Optional[str] = Form(None),
    glucose: Optional[str] = Form(None),
    # ── Image branch (optional) ──────────────────────────────────────────────
    ecg: Optional[UploadFile] = File(None),
):
    fields = {
        "male": male, "age": age, "education": education,
        "currentSmoker": currentSmoker, "cigsPerDay": cigsPerDay, "BPMeds": BPMeds,
        "prevalentStroke": prevalentStroke, "prevalentHyp": prevalentHyp,
        "diabetes": diabetes, "totChol": totChol, "sysBP": sysBP, "diaBP": diaBP,
        "BMI": BMI, "heartRate": heartRate, "glucose": glucose,
    }

    has_tabular = any(_has_value(fields[f]) for f in FEATURES)
    has_ecg = ecg is not None and bool(ecg.filename)

    if not has_tabular and not has_ecg:
        raise HTTPException(
            status_code=400,
            detail="Provide tabular patient data, an ECG image, or both.",
        )

    # Validate any provided tabular values (blanks are skipped → KNN-imputed).
    if has_tabular:
        errors = validate_tabular(fields)
        if errors:
            raise HTTPException(status_code=422, detail="; ".join(errors))

    branches: dict = {}

    # CPU-bound inference is offloaded to a threadpool so concurrent requests
    # don't block the event loop (torch/sklearn release the GIL).
    if has_tabular:
        try:
            from inference.framingham import predict_tabular
            branches["tabular"] = await run_in_threadpool(predict_tabular, fields)
        except FileNotFoundError as exc:
            raise HTTPException(status_code=503, detail=str(exc)) from exc

    if has_ecg:
        image_bytes = await ecg.read()
        try:
            from inference.ecg import predict_ecg
            branches["ecg"] = await run_in_threadpool(predict_ecg, image_bytes)
        except FileNotFoundError as exc:
            raise HTTPException(status_code=503, detail=str(exc)) from exc
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc

    # ── Fuse / select headline ───────────────────────────────────────────────
    if has_tabular and has_ecg:
        p = combine(branches["tabular"]["p_risk"], branches["ecg"]["p_risk"])
        mode, p_head = "multimodal", p
    elif has_tabular:
        mode, p_head = "tabular", branches["tabular"]["p_risk"]
    else:
        mode, p_head = "ecg", branches["ecg"]["p_risk"]

    return {
        "mode": mode,
        "risk_level": band_for(p_head),
        "p_risk": round(p_head, 4),
        "branches": branches,
    }


# ── Serve frontend ───────────────────────────────────────────────────────────
STATIC_DIR = Path(__file__).parent / "static"
STATIC_DIR.mkdir(exist_ok=True)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/")
def serve_frontend():
    return FileResponse(str(STATIC_DIR / "index.html"))