api/inference.py

from __future__ import annotations

import asyncio
import io
import os
import time

import torch
import torch.nn as nn
from PIL import Image

from data.transforms import build_transforms
from models.classifier import ChestAIClassifier
from models.uncertainty import mc_predict
from explainability.gradcam import ViTGradCAM
from report_generation.generator import (
    GROQ_AVAILABLE,
    RadiologyReportGenerator,
    generate_report_fallback,
)
from api.schemas import FindingResult, PredictionResponse
from data.dataset import CLASSES

THRESHOLDS = {
    "Atelectasis": 0.63, "Cardiomegaly": 0.74, "Effusion": 0.66,
    "Infiltration": 0.58, "Mass": 0.64, "Nodule": 0.58,
    "Pneumonia": 0.67, "Pneumothorax": 0.66, "Consolidation": 0.67,
    "Edema": 0.75, "Emphysema": 0.61, "Fibrosis": 0.60,
    "Pleural_Thickening": 0.61, "Hernia": 0.62,
}
UNCERTAINTY_THRESHOLD = 0.15


class InferencePipeline:
    _instance: "InferencePipeline | None" = None

    def __init__(self) -> None:
        self.model: ChestAIClassifier | None = None
        self.gradcam: ViTGradCAM | None = None
        self.report_gen: RadiologyReportGenerator | None = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.transform = build_transforms("val", image_size=224)
        self._lock = asyncio.Lock()
        self.model_version = "1.1.0"
        self.mc_samples = int(os.environ.get("MC_SAMPLES", "20"))

    async def load(self) -> None:
        from huggingface_hub import hf_hub_download

        model_repo = os.environ.get("MODEL_HUB_REPO", "Sowaiba01/chestai-model")
        checkpoint_path = os.environ.get("MODEL_CHECKPOINT", "")

        if not checkpoint_path and model_repo:
            print(f"[Pipeline] Downloading model from HF Hub: {model_repo}")
            checkpoint_path = hf_hub_download(
                repo_id=model_repo,
                filename="chestai_best.pt",
            )
        elif not checkpoint_path:
            raise RuntimeError(
                "Set MODEL_CHECKPOINT or MODEL_HUB_REPO env var."
            )

        print(f"[Pipeline] Loading checkpoint: {checkpoint_path}")
        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        state_dict = checkpoint["model_state_dict"]

        # The published checkpoint's classification head does not necessarily
        # match the default architecture in classifier.py (the shipped weights
        # use a 512-wide intermediate layer, while the class default is 256).
        # Read the true dimensions straight from the checkpoint and rebuild the
        # head to match, so the same code loads either variant. This block is
        # load-bearing: removing it produces a size-mismatch RuntimeError at
        # startup and the Space fails to boot.
        head_key = "head.2.weight"
        if head_key in state_dict:
            intermediate_dim = state_dict[head_key].shape[0]
            embed_dim = state_dict[head_key].shape[1]
        else:
            intermediate_dim = 512
            embed_dim = 512
        num_classes = len(CLASSES)
        dropout_rate = 0.3

        print(f"[Pipeline] Detected head: Linear({embed_dim} → {intermediate_dim} → {num_classes})")

        self.model = ChestAIClassifier()
        self.model.head = nn.Sequential(
            nn.LayerNorm(embed_dim),
            nn.Dropout(dropout_rate),
            nn.Linear(embed_dim, intermediate_dim),
            nn.GELU(),
            nn.Dropout(dropout_rate),
            nn.Linear(intermediate_dim, num_classes),
        ).to(self.device)

        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
        if missing:
            print(f"[Pipeline] Missing keys: {missing}")
        if unexpected:
            print(f"[Pipeline] Unexpected keys: {unexpected}")

        self.model.to(self.device)
        self.model.eval()

        self.gradcam = ViTGradCAM(self.model)

        groq_key = os.environ.get("GROQ_API_KEY")
        if groq_key and GROQ_AVAILABLE:
            self.report_gen = RadiologyReportGenerator(api_key=groq_key)
            print("[Pipeline] Groq report generator initialized.")
        elif groq_key and not GROQ_AVAILABLE:
            print("[Pipeline] GROQ_API_KEY is set but the 'groq' package is not "
                  "installed — using fallback report generator.")
        else:
            print("[Pipeline] GROQ_API_KEY not set — using fallback report generator.")

        print(f"[Pipeline] Model loaded on {self.device}.")

    async def _build_report(
        self,
        mean_probs,
        std_probs,
        patient_info: dict,
    ) -> str:
        """
        Generate the narrative report WITHOUT blocking the event loop.

        RadiologyReportGenerator.generate() is a synchronous network call to
        Groq (~1.2-1.8s). Awaiting it directly inside an async handler blocks
        the loop and serialises every concurrent request behind it, which is
        why p95 collapsed under load before this change. asyncio.to_thread
        moves it to the default executor.
        """
        def _sync() -> str:
            try:
                if self.report_gen:
                    return self.report_gen.generate(
                        probs=mean_probs.tolist(),
                        stds=std_probs.tolist(),
                        patient_info=patient_info or None,
                    )
                return generate_report_fallback(mean_probs.tolist(), std_probs.tolist())
            except Exception as e:  # noqa: BLE001 - never fail a scan on report error
                print(f"[Pipeline] Report generation failed: {e}")
                return generate_report_fallback(mean_probs.tolist(), std_probs.tolist())

        return await asyncio.to_thread(_sync)

    async def predict(
        self,
        image_bytes: bytes,
        patient_age: float | None = None,
        patient_gender: str | None = None,
        generate_report: bool = True,
        generate_gradcam: bool = True,
    ) -> PredictionResponse:
        """
        Full inference: preprocess → MC dropout → GradCAM → report.

        Set generate_report=False or generate_gradcam=False to skip those
        stages; both are optional and neither affects the pathology
        probabilities. See stage_timings_ms on the response for a per-stage
        breakdown.
        """
        if self.model is None:
            raise RuntimeError("Model not loaded. Call await pipeline.load() first.")

        t0 = time.perf_counter()
        stage: dict[str, float] = {}

        def _mark(name: str, start: float) -> float:
            now = time.perf_counter()
            stage[name] = round((now - start) * 1000, 1)
            return now

        pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        tensor = self.transform(pil_image).unsqueeze(0).to(self.device)
        t_pre = _mark("preprocess", t0)

        image_key = ViTGradCAM.image_key(image_bytes)

        # The model itself is not re-entrant (MC dropout toggles module state,
        # GradCAM mutates hook buffers), so GPU work stays under the lock.
        async with self._lock:
            mc_results = mc_predict(self.model, tensor, n_samples=self.mc_samples)
            mean_probs = mc_results["mean"][0].cpu().numpy()
            std_probs  = mc_results["std"][0].cpu().numpy()
            entropy    = float(mc_results["entropy"][0].item())
            t_mc = _mark("mc_dropout", t_pre)

            gradcam_classes = [
                cls for cls, p in zip(CLASSES, mean_probs)
                if p >= THRESHOLDS.get(cls, 0.5)
            ]

            if generate_gradcam and gradcam_classes:
                self.gradcam.generate_overlays(
                    tensor,
                    pil_image,
                    gradcam_classes,
                    image_key=image_key,
                )
            else:
                self.gradcam._last_overlays = {}
                gradcam_classes = [] if not generate_gradcam else gradcam_classes
            t_cam = _mark("gradcam", t_mc)

        findings = [
            FindingResult(
                name=cls,
                probability=float(p),
                uncertainty=float(s),
                present=float(p) >= THRESHOLDS.get(cls, 0.5),
                high_uncertainty=float(s) >= UNCERTAINTY_THRESHOLD,
            )
            for cls, p, s in zip(CLASSES, mean_probs, std_probs)
        ]

        patient_info = {}
        if patient_age is not None:
            patient_info["age"] = patient_age
        if patient_gender is not None:
            patient_info["gender"] = patient_gender

        if generate_report:
            report = await self._build_report(mean_probs, std_probs, patient_info)
        else:
            report = None
        _mark("report", t_cam)

        elapsed_ms = (time.perf_counter() - t0) * 1000

        return PredictionResponse(
            findings=findings,
            entropy=entropy,
            report=report,
            gradcam_available=bool(generate_gradcam and gradcam_classes),
            gradcam_classes=gradcam_classes,
            inference_time_ms=round(elapsed_ms, 1),
            stage_timings_ms=stage,
            model_version=self.model_version,
        )

    @property
    def is_loaded(self) -> bool:
        return self.model is not None


pipeline = InferencePipeline()