backend/main.py

import asyncio
import io
import json
import logging
import os
import random
import sys
import time
import uuid
from collections import deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import gradio as gr
import httpx
import numpy as np
from fastapi import (
    BackgroundTasks,
    FastAPI,
    File,
    HTTPException,
    Request,
    Response,
    UploadFile,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import PlainTextResponse
from PIL import Image
from prometheus_client import CONTENT_TYPE_LATEST, Counter, Histogram, generate_latest
from scipy.stats import ks_2samp
import matplotlib.cm as cm
import base64

# Heavy ML deps are optional — not installed on the Railway proxy node.
try:
    import torch
    import onnxruntime as ort
    from torchvision import transforms
    from torchvision.models import mobilenet_v3_small, efficientnet_b0
    from model_utils import CalibratedModel

    _TORCH_AVAILABLE = True
except ImportError:
    _TORCH_AVAILABLE = False


BASE_DIR = Path(__file__).resolve().parents[1]
LEGACY_MODEL_DIR = BASE_DIR / "models"
REALWORLD_MODEL_DIR = BASE_DIR / "models" / "realworld_efficientnet_b0"
CHESTMNIST_MODEL_DIR = BASE_DIR / "models" / "chestmnist_mobilenetv3"

PROFILE = os.getenv("PNEUMOOPS_PROFILE", "chestmnist").lower()
_DEFAULT_MODEL_DIR = {
    "realworld": REALWORLD_MODEL_DIR,
    "chestmnist": CHESTMNIST_MODEL_DIR,
}.get(PROFILE, LEGACY_MODEL_DIR)
MODEL_DIR = Path(os.getenv("PNEUMOOPS_MODEL_DIR", str(_DEFAULT_MODEL_DIR)))
REQUEST_LOG_HISTORY = deque(maxlen=20)

API_KEY = os.getenv("PNEUMOOPS_API_KEY")
ALLOWED_ORIGINS = [
    origin.strip()
    for origin in os.getenv("PNEUMOOPS_ALLOWED_ORIGINS", "*").split(",")
    if origin.strip()
]
TRAFFIC_WEIGHTS = {"pytorch": 60, "onnx": 40}
# When set, all model inference is forwarded to this HF Spaces URL instead of local models.
# Example: https://prakhar54-byte-pneumoops.hf.space
HF_SPACES_URL = os.getenv("HF_SPACES_URL", "").rstrip("/")
COLLECT_DATA = os.getenv("PNEUMOOPS_COLLECT_DATA", "true").lower() == "true"
COLLECT_DIR = BASE_DIR / "data" / "collected_images"
LOW_CONFIDENCE_THRESHOLD = float(
    os.getenv("PNEUMOOPS_LOW_CONFIDENCE_THRESHOLD", "0.60")
)
MIN_UPLOAD_EDGE = int(os.getenv("PNEUMOOPS_MIN_UPLOAD_EDGE", "96"))
MAX_CHANNEL_DELTA = float(os.getenv("PNEUMOOPS_MAX_CHANNEL_DELTA", "0.08"))
MIN_ASPECT_RATIO = float(os.getenv("PNEUMOOPS_MIN_ASPECT_RATIO", "0.6"))
MAX_ASPECT_RATIO = float(os.getenv("PNEUMOOPS_MAX_ASPECT_RATIO", "1.6"))

REQUEST_COUNTER = Counter(
    "pneumoops_requests_total",
    "Total inference requests served by PneumoOps.",
    ["model", "status"],
)
LATENCY_HISTOGRAM = Histogram(
    "pneumoops_inference_latency_ms",
    "Inference latency per model in milliseconds.",
    ["model"],
    buckets=(5, 10, 25, 50, 100, 250, 500, 1000),
)
DRIFT_COUNTER = Counter(
    "pneumoops_drift_alerts_total",
    "Number of drift alerts emitted by PneumoOps.",
    ["status"],
)
DISEASE_PREDICTION_COUNTER = Counter(
    "pneumoops_disease_predictions_total",
    "Per-disease prediction counts for production monitoring.",
    ["disease", "model"],
)

logger = logging.getLogger("pneumoops")
if not logger.handlers:
    logging.basicConfig(level=logging.INFO, format="%(message)s")


def resolve_path(candidates: list[Path]) -> Path | None:
    for candidate in candidates:
        if candidate.exists():
            return candidate
    return None


def load_json(path: Path | None, fallback: dict | None = None) -> dict:
    fallback_dict = {} if fallback is None else fallback
    if path and path.exists():
        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            logger.warning(f"Failed to parse JSON file {path}, using fallback.")
            return fallback_dict
    return fallback_dict


def resolve_runtime_paths(model_dir: Path) -> dict[str, Path | None]:
    checkpoint_path = resolve_path(
        [
            model_dir / "mobilenetv3_chestmnist.pth",  # chestmnist profile
            model_dir / "realworld_efficientnet_b0.pth",
            model_dir / "pneumo_model.pth",
        ]
    )
    onnx_report_path = resolve_path(
        [
            model_dir / "onnx_export_report.json",
            LEGACY_MODEL_DIR / "onnx_export_report.json",
        ]
    )
    onnx_report = load_json(onnx_report_path)

    base_onnx = onnx_report.get("base_onnx")
    optimized_onnx = onnx_report.get("optimized_onnx")
    serving_onnx = onnx_report.get("serving_onnx")

    onnx_candidates = []
    if serving_onnx:
        onnx_candidates.append(model_dir / serving_onnx)
    if optimized_onnx:
        onnx_candidates.append(model_dir / optimized_onnx)
    if base_onnx:
        onnx_candidates.append(model_dir / base_onnx)
    onnx_candidates.extend(
        [
            # chestmnist profile
            model_dir / "mobilenetv3_chestmnist.onnx",
            # realworld profile
            model_dir / "realworld_efficientnet_b0_quantized.onnx",
            model_dir / "realworld_efficientnet_b0_optimized.onnx",
            model_dir / "realworld_efficientnet_b0.onnx",
            model_dir / "pneumo_model_quantized.onnx",
            model_dir / "pneumo_model_optimized.onnx",
            model_dir / "pneumo_model.onnx",
        ]
    )

    return {
        "checkpoint": checkpoint_path,
        "onnx": resolve_path(onnx_candidates),
        "training_metrics": resolve_path(
            [
                model_dir / "training_metrics.json",
                LEGACY_MODEL_DIR / "training_metrics.json",
            ]
        ),
        "baseline_stats": resolve_path(
            [
                model_dir / "baseline_stats.json",
                LEGACY_MODEL_DIR / "baseline_stats.json",
            ]
        ),
        "onnx_export_report": onnx_report_path,
    }


def download_models_if_missing() -> None:
    """Download model weights from HF Hub at startup if not present locally.
    Only runs on the HF Spaces inference node (HF_SPACES_URL not set)."""
    if HF_SPACES_URL:
        return  # proxy node — no local models needed
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        return

    hf_model_repo = os.getenv("HF_MODEL_REPO", "Prakhar54-byte/pneumoops-chestmnist")
    hf_token = os.getenv("HF_TOKEN")
    files = ["mobilenetv3_chestmnist.pth", "mobilenetv3_chestmnist.onnx"]

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    for filename in files:
        dest = MODEL_DIR / filename
        if dest.exists():
            continue
        try:
            logger.info(f"Downloading {filename} from {hf_model_repo}...")
            hf_hub_download(
                repo_id=hf_model_repo,
                filename=filename,
                local_dir=str(MODEL_DIR),
                token=hf_token,
            )
            logger.info(f"Downloaded {filename} successfully.")
        except Exception as exc:
            logger.warning(f"Could not download {filename}: {exc}")


download_models_if_missing()
RUNTIME_PATHS = resolve_runtime_paths(MODEL_DIR)


def load_checkpoint_metadata(checkpoint_path: Path | None) -> dict[str, Any]:
    if checkpoint_path is None or not checkpoint_path.exists():
        # Proxy/backend-only deployments have no model weights but do have
        # training_metrics.json — use it so class names and thresholds are correct.
        tm_path = MODEL_DIR / "training_metrics.json"
        tm = load_json(tm_path)
        if tm.get("class_names"):
            class_names = tm["class_names"]
            n = len(class_names)
            return {
                "architecture": tm.get("architecture", "mobilenet_v3_small"),
                "class_names": class_names,
                "image_size": int(tm.get("image_size", 224)),
                "normalize_mean": [0.5, 0.5, 0.5],
                "normalize_std": [0.5, 0.5, 0.5],
                "thresholds": tm.get("thresholds", [0.5] * n),
                "logit_temperature": 1.0,
                "multi_label": bool(tm.get("multi_label", True)),
            }
        return {
            "architecture": "mobilenet_v3_small",
            "class_names": ["Normal", "Pneumonia"],
            "image_size": 224,
            "normalize_mean": [0.485, 0.456, 0.406],
            "normalize_std": [0.229, 0.224, 0.225],
            "thresholds": [0.5, 0.5],
            "logit_temperature": 1.0,
            "multi_label": False,
        }

    # If torch is not available, fall back to loading metadata from training_metrics.json
    if not _TORCH_AVAILABLE:
        tm_path = MODEL_DIR / "training_metrics.json"
        tm = load_json(tm_path)
        class_names = tm.get("class_names", ["Normal", "Pneumonia"])
        n = len(class_names)
        return {
            "architecture": tm.get("architecture", "mobilenet_v3_small"),
            "class_names": class_names,
            "image_size": int(tm.get("image_size", 224)),
            "normalize_mean": [0.485, 0.456, 0.406],
            "normalize_std": [0.229, 0.224, 0.225],
            "thresholds": tm.get("thresholds", [0.5] * n),
            "logit_temperature": 1.0,
            "multi_label": bool(tm.get("multi_label", True)),
        }

    payload = torch.load(checkpoint_path, map_location="cpu")
    if isinstance(payload, dict) and "model_state_dict" in payload:
        return {
            "state_dict": payload["model_state_dict"],
            "architecture": payload.get("architecture", "efficientnet_b0"),
            "class_names": payload.get("class_names", ["No Finding"]),
            "image_size": int(payload.get("image_size", 320)),
            "normalize_mean": payload.get("normalize_mean", [0.485, 0.456, 0.406]),
            "normalize_std": payload.get("normalize_std", [0.229, 0.224, 0.225]),
            "thresholds": payload.get(
                "thresholds", [0.5] * len(payload.get("class_names", ["No Finding"]))
            ),
            "logit_temperature": float(payload.get("logit_temperature", 1.0)),
            "multi_label": bool(payload.get("multi_label", True)),
        }

    # Plain state_dict (e.g. from train_chestmnist.py) — pull metadata from training_metrics.json
    tm_path = MODEL_DIR / "training_metrics.json"
    tm = load_json(tm_path)
    class_names = tm.get("class_names", ["Normal", "Pneumonia"])
    n = len(class_names)
    return {
        "state_dict": payload,
        "architecture": tm.get("architecture", "mobilenet_v3_small"),
        "class_names": class_names,
        "image_size": int(tm.get("image_size", 224)),
        "normalize_mean": [0.5, 0.5, 0.5],
        "normalize_std": [0.5, 0.5, 0.5],
        "thresholds": tm.get("thresholds", [0.5] * n),
        "logit_temperature": 1.0,
        "multi_label": bool(tm.get("multi_label", True)),
    }


MODEL_METADATA = load_checkpoint_metadata(RUNTIME_PATHS["checkpoint"])
CLASS_NAMES = MODEL_METADATA["class_names"]
MULTI_LABEL = bool(MODEL_METADATA["multi_label"])
IMAGE_SIZE = int(MODEL_METADATA["image_size"])
THRESHOLDS = np.asarray(MODEL_METADATA["thresholds"], dtype=np.float32)
LOGIT_TEMPERATURE = float(MODEL_METADATA.get("logit_temperature", 1.0))

TRANSFORM = (
    transforms.Compose(
        [
            transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
            transforms.Grayscale(num_output_channels=3),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=MODEL_METADATA["normalize_mean"],
                std=MODEL_METADATA["normalize_std"],
            ),
        ]
    )
    if _TORCH_AVAILABLE
    else None
)

BASELINE_STATS = load_json(
    RUNTIME_PATHS["baseline_stats"],
    fallback={
        "pixel_mean_mean": 0.5,
        "pixel_mean_std": 0.1,
        "pixel_std_mean": 0.2,
        "pixel_std_std": 0.05,
        "histogram_bins": 32,
        "histogram_mean": [1.0 / 32.0] * 32,
        "pixel_reference_sample": [],
        "drift_threshold": 1.2,
        "drift_ks_pvalue_threshold": 0.05,
        "class_names": CLASS_NAMES,
    },
)


def build_model() -> Any:
    checkpoint_path = RUNTIME_PATHS["checkpoint"]
    if checkpoint_path is None or not checkpoint_path.exists():
        return None

    architecture = MODEL_METADATA["architecture"]
    num_outputs = len(CLASS_NAMES)
    if architecture == "efficientnet_b0":
        base_model = efficientnet_b0(weights=None)
        base_model.classifier[1] = torch.nn.Linear(
            base_model.classifier[1].in_features, num_outputs
        )
    else:
        base_model = mobilenet_v3_small(weights=None)
        base_model.classifier[3] = torch.nn.Linear(
            base_model.classifier[3].in_features, num_outputs
        )

    base_model.load_state_dict(MODEL_METADATA["state_dict"])
    model = CalibratedModel(base_model, LOGIT_TEMPERATURE)
    model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
    model.eval()
    return model


def build_onnx_session():
    onnx_path = RUNTIME_PATHS["onnx"]
    if onnx_path is None or not onnx_path.exists():
        return None, None
    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    return session, onnx_path.name


DEVICE = (
    torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if _TORCH_AVAILABLE
    else None
)
PYTORCH_MODEL = build_model() if _TORCH_AVAILABLE else None
ONNX_SESSION, ACTIVE_ONNX_MODEL_NAME = (
    build_onnx_session() if _TORCH_AVAILABLE else (None, None)
)

app = FastAPI(
    title="PneumoOps API",
    description="Production-style MLOps pipeline for medical image classification with A/B routing and drift monitoring.",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    if API_KEY and request.url.path not in {"/health", "/metrics", "/"}:
        provided = request.headers.get("x-api-key")
        if provided != API_KEY:
            return Response(content="Unauthorized", status_code=401)
    return await call_next(request)


def load_image_from_upload(upload: UploadFile) -> Image.Image:
    try:
        content = upload.file.read()
        return Image.open(io.BytesIO(content)).convert("RGB")
    except Exception as exc:
        raise HTTPException(
            status_code=400, detail="Uploaded file is not a valid image."
        ) from exc


def summarize_image(image: Image.Image) -> dict[str, Any]:
    gray = np.asarray(image.convert("L"), dtype=np.float32) / 255.0
    width, height = image.size
    rgb = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
    channel_delta = (
        float(
            np.mean(np.abs(rgb[:, :, 0] - rgb[:, :, 1]))
            + np.mean(np.abs(rgb[:, :, 1] - rgb[:, :, 2]))
            + np.mean(np.abs(rgb[:, :, 0] - rgb[:, :, 2]))
        )
        / 3.0
    )
    return {
        "width": width,
        "height": height,
        "aspect_ratio": round(width / max(height, 1), 4),
        "pixel_mean": round(float(np.mean(gray)), 6),
        "pixel_std": round(float(np.std(gray)), 6),
        "pixel_min": round(float(np.min(gray)), 6),
        "pixel_max": round(float(np.max(gray)), 6),
        "channel_delta": round(channel_delta, 6),
    }


def validate_image(image: Image.Image, summary: dict[str, Any]) -> None:
    if min(summary["width"], summary["height"]) < MIN_UPLOAD_EDGE:
        raise HTTPException(
            status_code=400,
            detail=f"Input is too small for robust X-ray screening. Minimum edge is {MIN_UPLOAD_EDGE}px.",
        )
    if not (MIN_ASPECT_RATIO <= summary["aspect_ratio"] <= MAX_ASPECT_RATIO):
        raise HTTPException(
            status_code=400,
            detail="Input aspect ratio is outside the expected chest X-ray range.",
        )
    if summary["channel_delta"] > MAX_CHANNEL_DELTA:
        raise HTTPException(
            status_code=400,
            detail="Input appears to be a color image instead of a grayscale-style radiograph.",
        )


def calculate_drift(image: Image.Image) -> dict[str, Any]:
    gray = np.asarray(image.convert("L"), dtype=np.float32) / 255.0
    image_mean = float(np.mean(gray))
    image_std = float(np.std(gray))
    hist_bins = int(BASELINE_STATS.get("histogram_bins", 32))
    hist, _ = np.histogram(gray, bins=hist_bins, range=(0.0, 1.0), density=True)
    baseline_hist = np.asarray(
        BASELINE_STATS.get("histogram_mean", [1.0 / hist_bins] * hist_bins),
        dtype=np.float32,
    )

    mean_z = abs(image_mean - BASELINE_STATS.get("pixel_mean_mean", 0.5)) / max(
        BASELINE_STATS.get("pixel_mean_std", 0.1), 1e-6
    )
    std_z = abs(image_std - BASELINE_STATS.get("pixel_std_mean", 0.2)) / max(
        BASELINE_STATS.get("pixel_std_std", 0.05), 1e-6
    )
    hist_distance = float(np.mean(np.abs(hist - baseline_hist)))
    drift_score = round(0.35 * mean_z + 0.35 * std_z + 0.30 * hist_distance, 6)

    reference_sample = np.asarray(
        BASELINE_STATS.get("pixel_reference_sample", []), dtype=np.float32
    )
    incoming_sample = gray.reshape(-1)
    if reference_sample.size > 0:
        sample_size = min(len(incoming_sample), len(reference_sample), 4096)
        incoming_idx = np.random.choice(
            len(incoming_sample), size=sample_size, replace=False
        )
        reference_idx = np.random.choice(
            len(reference_sample), size=sample_size, replace=False
        )
        _, ks_pvalue = ks_2samp(
            incoming_sample[incoming_idx], reference_sample[reference_idx]
        )
        ks_pvalue = float(ks_pvalue)
    else:
        ks_pvalue = 1.0

    drift_detected = ks_pvalue < float(
        BASELINE_STATS.get("drift_ks_pvalue_threshold", 0.05)
    ) or drift_score > float(BASELINE_STATS.get("drift_threshold", 1.2))
    return {
        "drift_alert": "DRIFT_DETECTED" if drift_detected else "NORMAL",
        "drift_score": drift_score,
        "ks_pvalue": round(ks_pvalue, 6),
        "mean_z": round(float(mean_z), 6),
        "std_z": round(float(std_z), 6),
        "histogram_distance": round(hist_distance, 6),
    }


def postprocess_probabilities(probabilities: np.ndarray) -> dict[str, Any]:
    probabilities = probabilities.astype(np.float32)
    if MULTI_LABEL:
        predicted_indices = [
            index
            for index, value in enumerate(probabilities)
            if value >= THRESHOLDS[index]
        ]
        if not predicted_indices:
            predicted_indices = [int(np.argmax(probabilities))]
        predicted_labels = [CLASS_NAMES[index] for index in predicted_indices]
    else:
        predicted_index = int(np.argmax(probabilities))
        predicted_indices = [predicted_index]
        predicted_labels = [CLASS_NAMES[predicted_index]]

    all_predictions = [
        {
            "label": CLASS_NAMES[index],
            "confidence": round(float(probabilities[index]) * 100, 2),
            "threshold": round(float(THRESHOLDS[index]) * 100, 2)
            if index < len(THRESHOLDS)
            else 50.0,
            "detected": index in predicted_indices,
        }
        for index in range(len(CLASS_NAMES))
    ]
    sorted_pairs = sorted(
        all_predictions, key=lambda item: item["confidence"], reverse=True
    )
    top_confidence = sorted_pairs[0]["confidence"] if sorted_pairs else 0.0

    return {
        "predicted_labels": predicted_labels,
        "all_predictions": all_predictions, "summary": "Clinical Review Recommended: No findings reached the 30% threshold, but observation is advised for top results.",
        "top_predictions": sorted_pairs[: min(5, len(sorted_pairs))],
        "max_confidence": top_confidence,
        "low_confidence": top_confidence < (LOW_CONFIDENCE_THRESHOLD * 100.0),
        "inconclusive_scan": top_confidence < 10.0,
    }


class CAMHook:
    """Hook to capture gradients and activations for Grad-CAM."""

    def __init__(self, module):
        self.hook_f = module.register_forward_hook(self.hook_fn_fwd)
        # Score-CAM only needs forward hooks
        self.features = None
        

    def hook_fn_fwd(self, module, input, output):
        self.features = output.detach()

    def remove(self):
        self.hook_f.remove()

def generate_cam_overlay(image: Image.Image, cam_tensor: torch.Tensor) -> str:
    try:
        cam = cam_tensor.detach().cpu().numpy()
        cam = np.maximum(cam, 0)
        cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam) + 1e-12)
        w, h = image.size
        cam_img = Image.fromarray(np.uint8(255 * cam)).resize((w, h), Image.Resampling.LANCZOS)
        colormap = cm.get_cmap('jet')(np.array(cam_img) / 255.0)[:, :, :3]
        heatmap = np.uint8(255 * colormap)
        overlay = np.uint8(0.5 * np.array(image.convert('RGB')) + 0.5 * heatmap)
        buf = io.BytesIO()
        Image.fromarray(overlay).save(buf, format='PNG')
        return base64.b64encode(buf.getvalue()).decode('utf-8')
    except Exception as e:
        buf = io.BytesIO()
        image.save(buf, format='PNG')
        return base64.b64encode(buf.getvalue()).decode('utf-8')


def run_pytorch_inference(image: Image.Image) -> dict[str, Any]:
    if PYTORCH_MODEL is None:
        raise RuntimeError("PyTorch checkpoint is missing.")

    tensor = TRANSFORM(image).unsqueeze(0).to(DEVICE)
    tensor.requires_grad_(True)

    target_layer = None
    if hasattr(PYTORCH_MODEL.base_model, "features"):
        features = PYTORCH_MODEL.base_model.features
        def find_last_conv(m):
            last_conv = None
            for name, module in m.named_modules():
                if isinstance(module, torch.nn.Conv2d):
                    last_conv = module
            return last_conv
        target_layer = find_last_conv(features)
        if target_layer is None:
            target_layer = features[-1]
            logger.info(f"Using features[-1] as fallback target layer: {type(target_layer).__name__}")
        else:
            logger.info(f"Found deep Conv2d target layer: {type(target_layer).__name__}")
    else:
        logger.warning("No features attribute found on model for CAM")

    cam_hook = None
    if target_layer is not None:
        try:
            cam_hook = CAMHook(target_layer)
            logger.info(f"CAM hook registered on {type(target_layer).__name__}")
        except Exception as e:
            logger.error(f"Failed to register CAM hook: {e}")
    else:
        logger.warning("No target layer found for CAM")

    if torch.cuda.is_available():
        torch.cuda.synchronize()
    start = time.perf_counter()

    with torch.set_grad_enabled(True):
        logits = PYTORCH_MODEL(tensor)
        probs_tensor = (torch.sigmoid(logits).squeeze(0) if MULTI_LABEL else torch.softmax(logits, dim=1).squeeze(0))
        probabilities = probs_tensor.detach().cpu().numpy()
        class_idx = int(np.argmax(probabilities))
        
        cam_b64 = None
        if cam_hook:
            try:
                logger.info(f"Computing CAM for class index {class_idx}")
                PYTORCH_MODEL.zero_grad()
                
                # SCORE-CAM Implementation
                with torch.no_grad():
                    activations = cam_hook.features[0]  # [C, H, W]
                    num_channels = activations.shape[0]
            
                    # Upsample activations to match model input (224x224)
                    upsampled = torch.nn.functional.interpolate(
                        activations.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False
                    ).squeeze(0)  # [C, 224, 224]
            
                    # Normalize masks
                    upsampled = (upsampled - upsampled.min()) / (upsampled.max() - upsampled.min() + 1e-8)
            
                    # Batch process masked images to get scores
                    # We use a representative subset of channels if there are too many (e.g. 576) to keep latency low
                    # For MobileNetV3 small, we'll take top 128 channels or all if less
                    step = max(1, num_channels // 128)
                    subset_indices = list(range(0, num_channels, step))
            
                    masked_images = tensor * upsampled[subset_indices].unsqueeze(1) # [N, 3, 224, 224]
            
                    # Get scores for the target class
                    scores = PYTORCH_MODEL(masked_images)[:, class_idx]
                    scores = torch.nn.functional.softmax(scores, dim=0)
            
                    # Weighted sum of activations
                    cam_tensor = torch.sum(activations[subset_indices] * scores.view(-1, 1, 1), dim=0)
                    cam_b64 = generate_cam_overlay(image, cam_tensor)
                    logger.info("CAM overlay generated successfully")
            except Exception as e:
                logger.error(f"CAM computation failed: {e}")
            finally:
                cam_hook.remove()

    if torch.cuda.is_available():
        torch.cuda.synchronize()
    latency_ms = round((time.perf_counter() - start) * 1000, 2)

    return {
        "model_key": "pytorch",
        "model_used": "Baseline PyTorch",
        "latency_ms": latency_ms,
        "probabilities": probabilities.tolist(),
        "cam_b64": cam_b64,
        **postprocess_probabilities(probabilities),
    }


def run_onnx_inference(image: Image.Image) -> dict[str, Any]:
    if ONNX_SESSION is None:
        raise RuntimeError("ONNX artifact is missing.")

    tensor = TRANSFORM(image).unsqueeze(0).numpy().astype(np.float32)
    input_name = ONNX_SESSION.get_inputs()[0].name
    start = time.perf_counter()
    outputs = ONNX_SESSION.run(None, {input_name: tensor})
    latency_ms = round((time.perf_counter() - start) * 1000, 2)
    logits = torch.from_numpy(outputs[0]).squeeze(0)
    probabilities = (
        torch.sigmoid(logits).numpy()
        if MULTI_LABEL
        else torch.softmax(logits, dim=0).numpy()
    )
    return {
        "model_key": "onnx",
        "model_used": "Optimized ONNX",
        "latency_ms": latency_ms,
        "probabilities": probabilities.tolist(),
        **postprocess_probabilities(probabilities),
    }


async def _run_local_inference(image: Image.Image) -> dict[str, Any]:
    """Run both models in-process — used when model weights are available locally."""

    async def safe_call(model_name: str, fn):
        try:
            result = await asyncio.to_thread(fn, image)
            LATENCY_HISTOGRAM.labels(model=model_name).observe(result["latency_ms"])
            return result
        except Exception as exc:
            return {"model_key": model_name, "error": str(exc)}

    pytorch_result, onnx_result = await asyncio.gather(
        safe_call("pytorch", run_pytorch_inference),
        safe_call("onnx", run_onnx_inference),
    )
    return {"pytorch": pytorch_result, "onnx": onnx_result}


async def _call_hf_infer(image: Image.Image) -> dict[str, Any]:
    """Forward dual-model inference to the HF Spaces node via /infer."""
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    buf.seek(0)
    async with httpx.AsyncClient(timeout=90.0) as client:
        resp = await client.post(
            f"{HF_SPACES_URL}/infer",
            files={"file": ("xray.png", buf.read(), "image/png")},
        )
        resp.raise_for_status()
    data = resp.json()

    result: dict[str, Any] = {}
    for key in ("pytorch", "onnx"):
        raw = data.get(key) or {}
        if "error" in raw:
            result[key] = {"model_key": key, "error": raw["error"]}
        else:
            probs = np.asarray(raw.get("probabilities", []), dtype=np.float32)
            latency = float(raw.get("latency_ms") or 0.0)
            LATENCY_HISTOGRAM.labels(model=key).observe(latency)
            result[key] = {
                "model_key": key,
                "model_used": "Baseline PyTorch"
                if key == "pytorch"
                else "Optimized ONNX",
                "latency_ms": latency,
                "probabilities": probs.tolist(),
                **postprocess_probabilities(probs),
            }
    return result


async def benchmark_both_models(image: Image.Image) -> dict[str, Any]:
    """Route dual-model inference to HF Spaces (proxy mode) or local models."""
    if HF_SPACES_URL:
        return await _call_hf_infer(image)
    return await _run_local_inference(image)


def build_recommendation(
    selected_result: dict[str, Any],
    drift_result: dict[str, Any],
    dual_results: dict[str, Any],
) -> str:
    if drift_result["drift_alert"] == "DRIFT_DETECTED":
        return "Input distribution differs from the stored training baseline. Manual review is recommended before trusting this result."
    if selected_result["low_confidence"]:
        return "Prediction confidence is below the review threshold. Treat this as low confidence and escalate for human review."

    other_key = "onnx" if selected_result["model_key"] == "pytorch" else "pytorch"
    other_result = dual_results.get(other_key, {})
    if (
        other_result.get("predicted_labels")
        and other_result["predicted_labels"] != selected_result["predicted_labels"]
    ):
        return "The two serving paths disagree on the predicted findings. Use this as a monitoring alert and fall back to manual review."

    return "Use this output as a triage aid only. PneumoOps monitors latency and drift, but it is not a clinical decision-maker."


def append_history(entry: dict[str, Any]) -> None:
    REQUEST_LOG_HISTORY.append(entry)


def emit_structured_log(payload: dict[str, Any]) -> None:
    logger.info(json.dumps(payload, ensure_ascii=True))


def save_prediction_data(image: Image.Image, payload: dict[str, Any]) -> None:
    """Save anonymized image and prediction data for model fine-tuning."""
    if not COLLECT_DATA:
        return
    try:
        COLLECT_DIR.mkdir(parents=True, exist_ok=True)
        record_id = str(uuid.uuid4())

        # Save image
        img_path = COLLECT_DIR / f"{record_id}.png"
        image.save(img_path, format="PNG")

        # Save prediction payload
        json_path = COLLECT_DIR / f"{record_id}.json"
        json_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    except Exception as e:
        logger.error(f"Failed to save collection data: {e}")


@app.get("/health")
def health():
    return {
        "status": "ok",
        "deployment_mode": "proxy" if HF_SPACES_URL else "local",
        "hf_spaces_url": HF_SPACES_URL or None,
        "profile": PROFILE,
        "model_dir": str(MODEL_DIR),
        "pytorch_model_loaded": PYTORCH_MODEL is not None,
        "onnx_model_loaded": ONNX_SESSION is not None,
        "active_onnx_model": ACTIVE_ONNX_MODEL_NAME,
        "class_count": len(CLASS_NAMES),
        "class_names": CLASS_NAMES,
        "multi_label": MULTI_LABEL,
        "logit_temperature": LOGIT_TEMPERATURE,
        "training_metrics": load_json(RUNTIME_PATHS["training_metrics"]),
        "onnx_export_report": load_json(RUNTIME_PATHS["onnx_export_report"]),
        "recent_requests": list(REQUEST_LOG_HISTORY),
    }


@app.get("/metrics")
def metrics():
    return PlainTextResponse(generate_latest(), media_type=CONTENT_TYPE_LATEST)


@app.get("/history")
def history():
    return {"recent_requests": list(REQUEST_LOG_HISTORY)}


@app.get("/metrics/class-rates")
def class_prediction_rates():
    """Return per-class prediction rates from the rolling request history."""
    history_list = list(REQUEST_LOG_HISTORY)
    total = max(len(history_list), 1)
    rates: dict[str, float] = {label: 0.0 for label in CLASS_NAMES}
    avg_confidence: dict[str, list] = {label: [] for label in CLASS_NAMES}

    for entry in history_list:
        per_class = entry.get("per_class_predictions", {})
        per_class_probs = entry.get("per_class_probabilities", {})
        for label in CLASS_NAMES:
            if per_class.get(label, False):
                rates[label] = rates[label] + 1
            if label in per_class_probs:
                avg_confidence[label].append(per_class_probs[label])

    return {
        "window_size": total,
        "per_class_prediction_rate": {
            label: round(count / total, 4) for label, count in rates.items()
        },
        "per_class_avg_confidence": {
            label: round(float(sum(vals) / len(vals)), 4) if vals else None
            for label, vals in avg_confidence.items()
        },
        "drift_rate": round(
            sum(1 for e in history_list if e.get("drift_alert") == "DRIFT_DETECTED")
            / total,
            4,
        ),
    }


@app.get("/metrics/calibration")
def calibration_summary():
    """Return per-class calibration (Brier scores) from training metrics."""
    tm = load_json(RUNTIME_PATHS["training_metrics"])
    return {
        "test_macro_brier": tm.get("test_macro_brier"),
        "test_macro_auprc": tm.get("test_macro_auprc"),
        "test_macro_roc_auc": tm.get("test_macro_roc_auc"),
        "per_class_brier": tm.get("per_class_brier", {}),
        "per_class_auprc": tm.get("per_class_auprc", {}),
        "per_class_roc_auc": tm.get("per_class_roc_auc", {}),
        "per_class_recall": tm.get("per_class_recall", {}),
        "threshold_details": tm.get("threshold_details", {}),
        "class_names": tm.get("class_names", CLASS_NAMES),
    }


@app.post("/predict")
async def predict(
    request: Request, background_tasks: BackgroundTasks, file: UploadFile = File(...)
):
    image = load_image_from_upload(file)
    input_summary = summarize_image(image)
    validate_image(image, input_summary)

    start = time.perf_counter()
    dual_results = await benchmark_both_models(image)
    selected_key = random.choices(
        ["pytorch", "onnx"],
        weights=[TRAFFIC_WEIGHTS["pytorch"], TRAFFIC_WEIGHTS["onnx"]],
        k=1,
    )[0]
    selected_result = dual_results[selected_key]
    warning_flags = []

    if "error" in selected_result:
        fallback_result = dual_results["pytorch"]
        if "error" in fallback_result:
            REQUEST_COUNTER.labels(model=selected_key, status="failure").inc()
            raise HTTPException(
                status_code=503,
                detail=f"Both inference backends failed: {dual_results}",
            )
        selected_result = fallback_result
        warning_flags.append(f"{selected_key.upper()} failed, fallback to PyTorch.")

    # Increment per-class disease prediction counters
    for label in selected_result["predicted_labels"]:
        DISEASE_PREDICTION_COUNTER.labels(
            disease=label, model=selected_result["model_key"]
        ).inc()

    drift_result = calculate_drift(image)
    DRIFT_COUNTER.labels(status=drift_result["drift_alert"]).inc()
    REQUEST_COUNTER.labels(model=selected_result["model_key"], status="success").inc()

    pytorch_latency = (
        dual_results["pytorch"].get("latency_ms")
        if "error" not in dual_results["pytorch"]
        else None
    )
    onnx_latency = (
        dual_results["onnx"].get("latency_ms")
        if "error" not in dual_results["onnx"]
        else None
    )
    latency_delta = None
    if pytorch_latency is not None and onnx_latency is not None:
        latency_delta = round(float(onnx_latency) - float(pytorch_latency), 2)

    timestamp = datetime.now(timezone.utc).isoformat()
    response_payload = {
        "timestamp": timestamp,
        "selected_model": selected_result["model_used"],
        "selected_arm": "A" if selected_result["model_key"] == "pytorch" else "B",
        "weighted_traffic_split": TRAFFIC_WEIGHTS,
        "predicted_labels": selected_result["predicted_labels"],
        "all_predictions": selected_result["all_predictions"],
        "top_predictions": selected_result["top_predictions"],
        "confidence": selected_result["max_confidence"],
        "low_confidence": selected_result["low_confidence"],
        "inconclusive_scan": selected_result["inconclusive_scan"],
        "confidence_review_threshold": LOW_CONFIDENCE_THRESHOLD * 100.0,
        "pytorch_latency_ms": pytorch_latency,
        "onnx_latency_ms": onnx_latency,
        "latency_delta_ms": latency_delta,
        "latency_ms": selected_result["latency_ms"],
        "drift": drift_result,
        "input_summary": input_summary,
        "warning_flags": warning_flags,
        "recommendation": build_recommendation(
            selected_result, drift_result, dual_results
        ),
        "active_onnx_model": ACTIVE_ONNX_MODEL_NAME,
        "recent_history": list(REQUEST_LOG_HISTORY),
        "cam_b64": dual_results["pytorch"].get("cam_b64"),
    }

    history_entry = {
        "timestamp": timestamp,
        "model": selected_result["model_used"],
        "latency_ms": selected_result["latency_ms"],
        "drift_alert": drift_result["drift_alert"],
        "drift_score": drift_result["drift_score"],
        "confidence": selected_result["max_confidence"],
        "labels": selected_result["predicted_labels"],
        "per_class_probabilities": {
            CLASS_NAMES[i]: round(float(selected_result["probabilities"][i]), 4)
            for i in range(len(CLASS_NAMES))
            if i < len(selected_result.get("probabilities", []))
        },
        "per_class_predictions": {
            label: (label in selected_result["predicted_labels"])
            for label in CLASS_NAMES
        },
    }
    append_history(history_entry)
    response_payload["recent_history"] = list(REQUEST_LOG_HISTORY)

    emit_structured_log(
        {
            "event": "predict",
            "timestamp": timestamp,
            "model_used": selected_result["model_used"],
            "latency_ms": selected_result["latency_ms"],
            "confidence": selected_result["max_confidence"],
            "drift_status": drift_result["drift_alert"],
            "pathologies": selected_result["predicted_labels"],
            "client": request.client.host if request.client else None,
        }
    )

    response_payload["request_latency_ms"] = round(
        (time.perf_counter() - start) * 1000, 2
    )

    # Trigger background save for fine-tuning
    background_tasks.add_task(save_prediction_data, image, response_payload)

    return response_payload


@app.post("/infer")
async def infer_raw(file: UploadFile = File(...)):
    """Raw dual-model inference for remote backend nodes.
    Returns probabilities from both PyTorch and ONNX arms with no side-effects
    (no metrics, no drift, no history). Only available on the HF Spaces inference node
    (i.e. when HF_SPACES_URL is not set)."""
    if HF_SPACES_URL:
        raise HTTPException(
            status_code=501,
            detail="This node is a proxy backend; /infer is only served by the inference node.",
        )
    image = load_image_from_upload(file)
    dual = await _run_local_inference(image)
    out: dict[str, Any] = {}
    for key in ("pytorch", "onnx"):
        r = dual.get(key, {})
        out[key] = (
            {
                "probabilities": r.get("probabilities", []),
                "latency_ms": r.get("latency_ms"),
            }
            if "error" not in r
            else {"error": r["error"]}
        )
    return {
        **out,
        "class_names": CLASS_NAMES,
        "thresholds": THRESHOLDS.tolist(),
        "multi_label": MULTI_LABEL,
    }


# ─── Mount Gradio UI into FastAPI (single-port for HF Spaces) ───────────────
# This allows the entire app (API + UI) to run on one port (7860).
# - FastAPI REST endpoints remain at /predict, /health, /metrics, etc.
# - Gradio UI is served at / (root)

try:
    # Dynamically add frontend dir to path so app.py can be imported
    _root_dir = str(BASE_DIR)
    if _root_dir not in sys.path:
        sys.path.insert(0, _root_dir)

    import gradio as gr
    from frontend.app import demo as gradio_demo  # the gr.Blocks() object

    from fastapi.responses import RedirectResponse

    @app.get("/")
    def redirect_to_ui():
        return RedirectResponse(url="/ui")

    # Mount Gradio at /ui; FastAPI routes take priority because they're
    # registered first via @app.get / @app.post decorators.
    app = gr.mount_gradio_app(app, gradio_demo, path="/ui")
    logger.info("Gradio UI mounted at /ui — full app on single port.")

except Exception as _e:
    logger.warning(f"Gradio mount skipped ({_e}). API-only mode active.")


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)