api/main.py

"""
api/main.py
───────────
FastAPI layer exposing the breast cancer classifier as an HTTP API.

Endpoints
─────────
  GET  /health                — Liveness check
  POST /predict               — Single image inference (histopathology)
  POST /predict/batch         — Multi-image batch inference
  POST /explain               — Prediction + LLM natural language report
  POST /explain/visual        — Prediction + Grad-CAM heatmap (base64 PNG)
  POST /chat                  — Conversational follow-up
  POST /mammogram/predict     — Mammogram inference (3-model ensemble)
  POST /mammogram/explain     — Mammogram + LLM explanation
  POST /mammogram/visual      — Mammogram + Grad-CAM + LLM explanation

Run locally
───────────
  uvicorn api.main:app --host 0.0.0.0 --port 8000 --reload

Env vars
────────
  WEIGHTS_PATH       model/weights.pth               (histopathology)
  MAMMO_WEIGHTS_DIR  Directory of ensemble members   (default: model)
  MAMMO_THRESHOLD    Ensemble decision threshold     (default: 0.5)
  DEVICE             "cuda" | "mps" | "cpu"          (default: auto)
  CONFIDENCE_THR     Float in [0,1]                  (default: 0.5)
  USE_LLM_MODEL      "true" | "false"                (default: false)
  LLM_MODEL_NAME     HuggingFace model id            (default: flan-t5-base)
  GRADCAM_ALPHA      Heatmap blend opacity 0–1       (default: 0.5)
"""

from __future__ import annotations

import base64
import io
import logging
import math
import os
import sys
from pathlib import Path
from typing import List, Optional

import torch
from fastapi import FastAPI, File, HTTPException, Query, UploadFile, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from PIL import Image

# Ensure project root is on PYTHONPATH when running from repo root
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from model.inference             import BreastCancerInferencePipeline
from model.mammogram_inference   import MammogramInferencePipeline
from model.mammogram_ensemble    import MammogramEnsemble
from explainability.gradcam            import GradCAM
from explainability.mammogram_gradcam  import MammogramGradCAM
from explainability.chat_pipeline   import DualModelChatPipeline, create_pipeline
from explainability.llm_explain import LLMExplainer, ChatEngine
from explainability.llm_chat import stream_chat, llm_available

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ── FastAPI app ───────────────────────────────────────────────────────────────
app = FastAPI(
    title="MedAI — Breast Cancer Analysis Platform",
    description=(
        "DenseNet-121 histopathology classifier and 3-model EfficientNet-B4 "
        "mammogram ensemble, with Grad-CAM explainability and LLM explanation. "
        "Research and educational use only. Not a standalone diagnostic tool."
    ),
    version="2.1.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)


# ── Module singletons ─────────────────────────────────────────────────────────
_pipeline:       BreastCancerInferencePipeline | None              = None
_mammo_pipeline: "MammogramEnsemble | MammogramInferencePipeline | None" = None
_gradcam:        GradCAM | None                                    = None
_llm_explainer:  LLMExplainer | None                               = None
_chat_engine:    ChatEngine | None                                 = None
_chat_pipeline:  DualModelChatPipeline | None                      = None


@app.on_event("startup")
def _load_pipeline() -> None:
    global _pipeline, _mammo_pipeline
    weights = os.getenv("WEIGHTS_PATH", "model/weights.pth")
    device  = os.getenv("DEVICE", None)
    thr     = float(os.getenv("CONFIDENCE_THR", "0.5"))

    weights_path = Path(weights) if Path(weights).exists() else None
    if weights_path is None:
        logger.warning(
            "weights.pth not found at '%s'. "
            "Running with ImageNet-pretrained backbone only.", weights,
        )

    _pipeline = BreastCancerInferencePipeline(
        weights_path=weights_path,
        device=device,
        confidence_threshold=thr,
    )
    logger.info("Histopathology pipeline ready on device: %s", _pipeline.device)

    # ── Mammogram: prefer the 3-model ensemble; fall back to single-view ──────
    mammo_dir = Path(os.getenv("MAMMO_WEIGHTS_DIR", "model"))
    mammo_thr = float(os.getenv("MAMMO_THRESHOLD", "0.5"))
    members   = [mammo_dir / f"model_s{s}.pth" for s in (42, 123, 999)]
    present   = [p for p in members if p.exists()]

    if present:
        _mammo_pipeline = MammogramEnsemble(
            weight_paths = present,
            device       = device,
            threshold    = mammo_thr,
        )
        logger.info(
            "Mammogram ENSEMBLE ready on %s (%d members, threshold=%.2f)",
            _mammo_pipeline.device, len(present), mammo_thr,
        )
    else:
        legacy = mammo_dir / "mammogram_weights.pth"
        _mammo_pipeline = MammogramInferencePipeline(
            weights_path = legacy if legacy.exists() else None,
            device       = device,
        )
        logger.warning(
            "Ensemble members not found in '%s'. Falling back to single-view "
            "pipeline (%s).", mammo_dir,
            "trained weights" if legacy.exists() else "ImageNet init",
        )

    logger.info("Grad-CAM and LLM explainer load on first /explain request.")


def _get_mammo_pipeline():
    if _mammo_pipeline is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Mammogram pipeline not ready.",
        )
    return _mammo_pipeline


def _get_pipeline() -> BreastCancerInferencePipeline:
    if _pipeline is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Model pipeline is not ready. Try again shortly.",
        )
    return _pipeline


def _get_gradcam() -> GradCAM:
    """Lazy-load GradCAM on first explainability request (histopathology)."""
    global _gradcam
    if _gradcam is None:
        pipeline = _get_pipeline()
        alpha = float(os.getenv("GRADCAM_ALPHA", "0.5"))
        logger.info("Initialising Grad-CAM module (device=%s, alpha=%.2f)...",
                    pipeline.device, alpha)
        _gradcam = GradCAM(
            model  = pipeline.model,
            device = str(pipeline.device),
            alpha  = alpha,
        )
        logger.info("Grad-CAM ready.")
    return _gradcam


def _get_llm() -> LLMExplainer:
    global _llm_explainer
    if _llm_explainer is None:
        use_llm    = os.getenv("USE_LLM_MODEL", "false").lower() == "true"
        model_name = os.getenv("LLM_MODEL_NAME", "google/flan-t5-base")
        logger.info(
            "Initialising LLM explainer (use_llm=%s, model=%s)...",
            use_llm, model_name if use_llm else "template-engine",
        )
        _llm_explainer = LLMExplainer(model_name=model_name, use_llm=use_llm)
        logger.info("LLM explainer ready (engine=%s).",
                    "flan-t5" if use_llm else "template")
    return _llm_explainer


def _get_chat_engine() -> ChatEngine:
    global _chat_engine
    if _chat_engine is None:
        _chat_engine = ChatEngine()
        logger.info("ChatEngine fallback ready.")
    return _chat_engine


def _get_chat_pipeline() -> "DualModelChatPipeline | None":
    global _chat_pipeline
    if _chat_pipeline is None:
        _chat_pipeline = create_pipeline()
        if _chat_pipeline:
            logger.info("Dual pipeline ready: BioMedLM + Llama 3.2 (Groq)")
        else:
            logger.info(
                "Dual pipeline not available — set GROQ_API_KEY + HF_TOKEN "
                "to enable BioMedLM + Llama 3.2. Using ChatEngine fallback."
            )
    return _chat_pipeline


# ── Response schemas ──────────────────────────────────────────────────────────

class PredictionResponse(BaseModel):
    prediction: str        = Field(..., examples=["malignant"])
    confidence: float      = Field(..., ge=0.0, le=1.0, examples=[0.875])
    logits:     list[float]= Field(
        ...,
        description="Raw pre-softmax scores [benign, malignant]",
        examples=[[-2.14, 3.87]],
    )
    disclaimer: str = Field(
        default=("Research / educational use only. "
                 "Not a standalone diagnostic tool.")
    )


class ExplainResponse(PredictionResponse):
    summary:    str = Field(..., description="Plain-language summary of the prediction")
    detail:     str = Field(..., description="Deeper explanation with confidence context")
    audience:   str = Field(..., description="Target audience the explanation is tailored for")
    engine:     str = Field(..., description="'flan-t5' if LLM ran, 'template' if fallback")


class VisualExplainResponse(ExplainResponse):
    overlay_b64:      str   = Field(..., description="Base64 PNG of Grad-CAM overlay")
    heatmap_b64:      str   = Field(..., description="Base64 PNG of raw Grad-CAM heatmap")
    spatial_summary:  str   = Field(..., description="Regions that drove the prediction")
    heatmap_mean:     float = Field(..., description="Mean activation value [0,1]")
    heatmap_max:      float = Field(..., description="Peak activation value [0,1]")


class HealthResponse(BaseModel):
    status:          str
    device:          str
    gradcam_loaded:  bool
    llm_loaded:      bool
    llm_engine:      str
    mammogram_mode:  str
    chat_llm:        str


# ── Helpers ───────────────────────────────────────────────────────────────────
ACCEPTED_MIME = {"image/jpeg", "image/png", "image/tiff", "image/bmp"}


def _validate_and_open(upload: UploadFile) -> Image.Image:
    if upload.content_type not in ACCEPTED_MIME:
        raise HTTPException(
            status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
            detail=f"Unsupported file type: {upload.content_type}. "
                   f"Accepted: {', '.join(ACCEPTED_MIME)}",
        )
    data = upload.file.read()
    try:
        return Image.open(io.BytesIO(data)).convert("RGB")
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=f"Could not decode image: {exc}",
        ) from exc


def _serialize(result: dict) -> dict:
    return {
        "prediction": result["prediction"],
        "confidence": result["confidence"],
        "logits":     result["logits"].squeeze().tolist(),
    }


def _pil_to_b64(img: Image.Image, fmt: str = "PNG") -> str:
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def _heatmap_to_pil(heatmap) -> Image.Image:
    import numpy as np
    arr = (heatmap * 255).astype(np.uint8)
    return Image.fromarray(arr, mode="L")


def _ensemble_logits(result: dict) -> list[float]:
    """
    The ensemble returns averaged probabilities, not raw logits. Synthesize
    log-probabilities so the response schema and the /chat context (which
    expects [benign, malignant] scores) stay backward-compatible.
    """
    mal = float(result.get("malignant_probability", result.get("confidence", 0.5)))
    ben = 1.0 - mal
    return [round(math.log(max(ben, 1e-6)), 6),
            round(math.log(max(mal, 1e-6)), 6)]


def _is_ensemble(pipeline) -> bool:
    return isinstance(pipeline, MammogramEnsemble)


class PatientRecord(BaseModel):
    name:             str = Field(default="", description="Patient name")
    age:              int = Field(default=0,  description="Patient age")
    sex:              str = Field(default="", description="Male / Female / Other")
    medical_history:  str = Field(default="", description="Relevant medical history")
    symptoms:         str = Field(default="", description="Current symptoms or concerns")
    previous_scans:   str = Field(default="", description="Previous imaging history")


class ChatRequest(BaseModel):
    message:         str
    audience:        str   = Field(default="clinician",
                                   description="clinician | researcher | patient")
    prediction:      str   = Field(default="")
    confidence:      float = Field(default=0.0, ge=0.0, le=1.0)
    logits:          list[float] = Field(default=[0.0, 0.0])
    spatial_summary: str   = Field(default="")
    history:         list  = Field(default_factory=list)
    patient:         PatientRecord = Field(default_factory=PatientRecord)


class ChatResponse(BaseModel):
    response:  str
    audience:  str


# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.get("/health", response_model=HealthResponse, tags=["Meta"])
def health() -> HealthResponse:
    pipeline = _get_pipeline()
    llm_engine = "not_loaded"
    if _llm_explainer is not None:
        llm_engine = "flan-t5" if _llm_explainer.use_llm else "template"

    mammo_mode = "not_loaded"
    if _mammo_pipeline is not None:
        mammo_mode = "ensemble" if _is_ensemble(_mammo_pipeline) else "single-view"

    import os
    avail, provider = llm_available()
    model = os.getenv("CHAT_MODEL") or ("openai/gpt-oss-120b" if provider == "groq" else provider)
    chat_llm = f"{provider}/{model}" if avail else f"{provider}/no_api_key"

    return HealthResponse(
        status         = "ok",
        device         = str(pipeline.device),
        gradcam_loaded = _gradcam is not None,
        llm_loaded     = _llm_explainer is not None,
        llm_engine     = llm_engine,
        mammogram_mode = mammo_mode,
        chat_llm       = chat_llm,
    )


@app.post("/predict", response_model=PredictionResponse, tags=["Inference"])
def predict(file: UploadFile = File(...)) -> PredictionResponse:
    pipeline = _get_pipeline()
    image    = _validate_and_open(file)
    try:
        result = pipeline.predict(image)
    except Exception as exc:
        logger.exception("Inference error")
        raise HTTPException(500, detail=f"Inference failed: {exc}") from exc
    return PredictionResponse(**_serialize(result))


@app.post("/predict/batch", response_model=List[PredictionResponse], tags=["Inference"])
def predict_batch(files: List[UploadFile] = File(...)) -> list[PredictionResponse]:
    if len(files) > 16:
        raise HTTPException(400, detail="Batch size exceeds maximum of 16 images.")
    pipeline = _get_pipeline()
    images   = [_validate_and_open(f) for f in files]
    try:
        results = pipeline.predict_batch(images)
    except Exception as exc:
        logger.exception("Batch inference error")
        raise HTTPException(500, detail=f"Batch inference failed: {exc}") from exc
    return [PredictionResponse(**_serialize(r)) for r in results]


@app.post("/explain", response_model=ExplainResponse, tags=["Explainability"])
def explain(
    file:     UploadFile = File(...),
    audience: str        = Query(default="clinician",
                                 enum=["clinician", "researcher", "patient"]),
) -> ExplainResponse:
    pipeline = _get_pipeline()
    llm      = _get_llm()
    image    = _validate_and_open(file)
    try:
        prediction = pipeline.predict(image)
    except Exception as exc:
        logger.exception("Inference error in /explain")
        raise HTTPException(500, detail=f"Inference failed: {exc}") from exc
    try:
        report = llm.explain(prediction, audience=audience)
    except Exception as exc:
        logger.exception("LLM explanation error in /explain")
        raise HTTPException(500, detail=f"Explanation generation failed: {exc}") from exc
    return ExplainResponse(
        **_serialize(prediction),
        summary  = report["summary"],
        detail   = report["detail"],
        audience = report["audience"],
        engine   = report["engine"],
    )


@app.post("/explain/visual", response_model=VisualExplainResponse, tags=["Explainability"])
def explain_visual(
    file:      UploadFile    = File(...),
    audience:  str           = Query(default="clinician",
                                     enum=["clinician", "researcher", "patient"]),
    class_idx: Optional[int] = Query(default=None, ge=0, le=1),
) -> VisualExplainResponse:
    llm   = _get_llm()
    cam   = _get_gradcam()
    image = _validate_and_open(file)
    try:
        cam_result = cam.explain(image, class_idx=class_idx)
    except Exception as exc:
        logger.exception("Grad-CAM error in /explain/visual")
        raise HTTPException(500, detail=f"Grad-CAM failed: {exc}") from exc
    try:
        report = llm.explain_with_gradcam(cam_result, audience=audience)
    except Exception as exc:
        logger.exception("LLM explanation error in /explain/visual")
        raise HTTPException(500, detail=f"Explanation generation failed: {exc}") from exc

    import numpy as np
    heatmap = cam_result["heatmap"]
    overlay = cam_result["overlay"]
    return VisualExplainResponse(
        **_serialize(cam_result),
        summary         = report["summary"],
        detail          = report["detail"],
        audience        = report["audience"],
        engine          = report["engine"],
        overlay_b64     = _pil_to_b64(overlay),
        heatmap_b64     = _pil_to_b64(_heatmap_to_pil(heatmap)),
        spatial_summary = LLMExplainer._summarise_heatmap(heatmap),
        heatmap_mean    = round(float(np.mean(heatmap)), 6),
        heatmap_max     = round(float(np.max(heatmap)), 6),
    )


@app.post("/chat", response_model=ChatResponse, tags=["Explainability"])
def chat(request: ChatRequest) -> ChatResponse:
    logits = request.logits if len(request.logits) >= 2 else [0.0, 0.0]
    patient_dict = {}
    if request.patient:
        patient_dict = {
            "name":            request.patient.name,
            "age":             request.patient.age,
            "sex":             request.patient.sex,
            "medical_history": request.patient.medical_history,
            "symptoms":        request.patient.symptoms,
            "previous_scans":  request.patient.previous_scans,
        }

    pipeline = _get_chat_pipeline()
    if pipeline:
        try:
            response_text = pipeline.respond(
                message         = request.message,
                audience        = request.audience,
                prediction      = request.prediction,
                confidence      = request.confidence,
                benign_logit    = logits[0],
                malignant_logit = logits[1],
                spatial_summary = request.spatial_summary,
                history         = request.history,
                patient         = patient_dict,
            )
            if response_text:
                return ChatResponse(response=response_text, audience=request.audience)
        except Exception as e:
            logger.warning("Dual pipeline error (%s) — falling back to ChatEngine.", e)

    _get_llm()
    engine = _get_chat_engine()
    response_text = engine.respond(
        message         = request.message,
        audience        = request.audience,
        prediction      = request.prediction,
        confidence      = request.confidence,
        benign_logit    = logits[0],
        malignant_logit = logits[1],
        spatial_summary = request.spatial_summary,
        history         = request.history,
        patient         = patient_dict,
    )
    return ChatResponse(response=response_text, audience=request.audience)


class LLMChatRequest(BaseModel):
    messages:        list  = Field(default_factory=list,
                                    description="[{role: 'user'|'assistant', content: str}, ...]")
    audience:        str   = Field(default="clinician")
    prediction:      str   = Field(default="")
    confidence:      float = Field(default=0.0, ge=0.0, le=1.0)
    logits:          list[float] = Field(default=[0.0, 0.0])
    spatial_summary: str   = Field(default="")
    birads:          str   = Field(default="")


@app.post("/chat/stream", tags=["Explainability"])
def chat_stream(req: LLMChatRequest):
    """Real LLM assistant with token streaming (Claude or GPT via env config)."""
    context = {
        "prediction":      req.prediction,
        "confidence":      req.confidence,
        "logits":          req.logits,
        "spatial_summary": req.spatial_summary,
        "birads":          req.birads,
    } if req.prediction else None

    def generate():
        try:
            for token in stream_chat(req.messages, context=context, audience=req.audience):
                yield token
        except Exception as e:  # noqa: BLE001
            logger.warning("LLM chat stream error: %s", e)
            yield f"\n[Assistant error: {e}]"

    return StreamingResponse(generate(), media_type="text/plain; charset=utf-8")

class MammogramResponse(BaseModel):
    prediction:  str
    confidence:  float
    logits:      list[float]
    birads:      str
    modality:    str = "mammogram"
    per_model:   Optional[list[float]] = Field(
        default=None,
        description="Per-member malignant probabilities (ensemble mode only)",
    )
    n_models:    Optional[int] = Field(
        default=None, description="Number of ensemble members averaged",
    )
    disclaimer:  str = (
        "AI-assisted mammogram analysis. Research use only. "
        "Not a substitute for radiologist review or clinical diagnosis."
    )


class MammogramExplainResponse(MammogramResponse):
    summary:    str
    detail:     str
    audience:   str
    engine:     str


class MammogramVisualResponse(MammogramExplainResponse):
    overlay_b64:     str
    heatmap_b64:     str
    spatial_summary: str
    heatmap_mean:    float
    heatmap_max:     float


def _run_mammo_predict(pipeline, image: Image.Image, tta: bool) -> dict:
    """Unified prediction → dict with prediction/confidence/logits/birads/etc."""
    if _is_ensemble(pipeline):
        result = pipeline.predict_tta(image) if tta else pipeline.predict(image)
        return {
            "prediction": result["prediction"],
            "confidence": result["confidence"],
            "logits":     _ensemble_logits(result),
            "birads":     result["birads"],
            "modality":   result.get("modality", "mammogram_ensemble"),
            "per_model":  result.get("per_model"),
            "n_models":   result.get("n_models"),
        }
    # Legacy single-view pipeline (returns logits tensor)
    result = pipeline.predict(image)
    return {
        "prediction": result["prediction"],
        "confidence": result["confidence"],
        "logits":     result["logits"].squeeze().tolist(),
        "birads":     result["birads"],
        "modality":   result.get("modality", "mammogram"),
        "per_model":  None,
        "n_models":   None,
    }


@app.post("/mammogram/predict",
          response_model=MammogramResponse,
          tags=["Mammogram"])
def mammogram_predict(
    file: UploadFile = File(...),
    tta:  bool       = Query(default=False,
                             description="Test-time augmentation (ensemble only) — "
                                         "slower, slightly more accurate"),
) -> MammogramResponse:
    """
    Classify a mammogram using the 3-model EfficientNet-B4 ensemble.
    Returns prediction, confidence, BI-RADS suggestion, and per-member probs.
    """
    pipeline = _get_mammo_pipeline()
    image    = _validate_and_open(file)
    try:
        result = _run_mammo_predict(pipeline, image, tta)
    except Exception as exc:
        logger.exception("Mammogram inference error")
        raise HTTPException(500, detail=f"Mammogram inference failed: {exc}") from exc
    return MammogramResponse(**result)


@app.post("/mammogram/explain",
          response_model=MammogramExplainResponse,
          tags=["Mammogram"])
def mammogram_explain(
    file:     UploadFile = File(...),
    audience: str        = Query(default="clinician",
                                 enum=["clinician", "researcher", "patient"]),
    tta:      bool       = Query(default=False),
) -> MammogramExplainResponse:
    """Classify a mammogram (ensemble) and generate an audience-specific explanation."""
    pipeline = _get_mammo_pipeline()
    llm      = _get_llm()
    image    = _validate_and_open(file)

    try:
        result = _run_mammo_predict(pipeline, image, tta)
    except Exception as exc:
        raise HTTPException(500, detail=f"Mammogram inference failed: {exc}") from exc

    explain_input = {
        "prediction": result["prediction"],
        "confidence": result["confidence"],
        "logits":     torch.tensor([result["logits"]]),
        "birads":     result["birads"],
    }
    try:
        report = llm.explain(explain_input, audience=audience, modality="mammogram")
    except Exception as exc:
        raise HTTPException(500, detail=f"Explanation failed: {exc}") from exc

    return MammogramExplainResponse(
        **result,
        summary  = report["summary"],
        detail   = report["detail"],
        audience = report["audience"],
        engine   = report["engine"],
    )


@app.post("/mammogram/visual",
          response_model=MammogramVisualResponse,
          tags=["Mammogram"])
def mammogram_visual(
    file:      UploadFile    = File(...),
    audience:  str           = Query(default="clinician",
                                     enum=["clinician", "researcher", "patient"]),
    class_idx: Optional[int] = Query(default=None, ge=0, le=1),
) -> MammogramVisualResponse:
    """
    Full mammogram analysis: ensemble prediction + Grad-CAM + LLM explanation.

    Grad-CAM is inherently a single-model technique, so the heatmap is computed
    on the strongest ensemble member (the most representative saliency map),
    while the prediction label/confidence uses the full ensemble average.
    """
    import numpy as np

    pipeline = _get_mammo_pipeline()
    llm      = _get_llm()
    image    = _validate_and_open(file)

    # Authoritative label/confidence from the full ensemble
    try:
        pred = _run_mammo_predict(pipeline, image, tta=False)
    except Exception as exc:
        raise HTTPException(500, detail=f"Mammogram inference failed: {exc}") from exc

    # Grad-CAM on the representative single model (EfficientNet-specific)
    cam_model = pipeline.cam_model if _is_ensemble(pipeline) else pipeline.model
    try:
        alpha = float(os.getenv("GRADCAM_ALPHA", "0.5"))
        mammo_cam = MammogramGradCAM(model=cam_model, device=str(pipeline.device), alpha=alpha)
        cam_result = mammo_cam.explain(image, class_idx=class_idx)
    except Exception as exc:
        logger.exception("Mammogram Grad-CAM error")
        raise HTTPException(
            500,
            detail=(f"Grad-CAM failed: {exc}. The /mammogram/predict endpoint "
                    f"still works without the heatmap."),
        ) from exc

    # Make the LLM explanation reflect the ENSEMBLE verdict (authoritative),
    # while the heatmap stays the representative member's saliency map.
    cam_result["prediction"] = pred["prediction"]
    cam_result["confidence"] = pred["confidence"]
    cam_result["logits"]     = torch.tensor([pred["logits"]])
    cam_result["birads"]     = pred["birads"]

    try:
        report = llm.explain_with_gradcam(cam_result, audience=audience, modality="mammogram")
    except Exception as exc:
        raise HTTPException(500, detail=f"Explanation failed: {exc}") from exc

    heatmap = cam_result["heatmap"]
    overlay = cam_result["overlay"]
    return MammogramVisualResponse(
        prediction      = pred["prediction"],
        confidence      = pred["confidence"],
        logits          = pred["logits"],
        birads          = pred["birads"],
        modality        = pred["modality"],
        per_model       = pred["per_model"],
        n_models        = pred["n_models"],
        summary         = report["summary"],
        detail          = report["detail"],
        audience        = report["audience"],
        engine          = report["engine"],
        overlay_b64     = _pil_to_b64(overlay),
        heatmap_b64     = _pil_to_b64(_heatmap_to_pil(heatmap)),
        spatial_summary = LLMExplainer._summarise_heatmap(heatmap),
        heatmap_mean    = round(float(np.mean(heatmap)), 6),
        heatmap_max     = round(float(np.max(heatmap)), 6),
    )