server.py

"""
RetViM - Real Inference Server v4.0
====================================
Uses ImprovedMedMamba (model.py) with the real .ckpt weights.
val_acc = 96.68 %  *  ViT-B/16 + 2 SS-Conv-SSM blocks  *  dim = 768

Run:  python server.py
Open: http://127.0.0.1:8000/
"""
from __future__ import annotations
import io, os, base64, time, json, hashlib, math, gc
import re
from pathlib import Path
from typing import Optional, List

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy.ndimage

from fastapi import FastAPI, File, UploadFile, Query, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse, Response, FileResponse
from fastapi.staticfiles import StaticFiles
import uvicorn

# Import the real model
from model import load_model, ImprovedMedMamba

# --------------------------------------------------------------------------- #
#  Config
# --------------------------------------------------------------------------- #
BASE_DIR   = Path(__file__).parent
HTML_FILE  = BASE_DIR / "retvim-final.html"
CKPT_PATH  = BASE_DIR / "improved-medmamba-epoch=19-val_acc=0.9668.ckpt"
DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
PORT       = 8000
HOST       = "127.0.0.1"
START_TIME = time.time()

CLASS_NAMES   = ["CNV", "DME", "DRUSEN", "NORMAL"]
SAMPLE_DIR    = BASE_DIR / "sample images"
SAMPLE_IMAGES = {
    "CNV":    SAMPLE_DIR / "CNV-103044-6.jpeg",
    "DME":    SAMPLE_DIR / "DME-30521-3.jpeg",
    "DRUSEN": SAMPLE_DIR / "DRUSEN-1786810-1.jpeg",
    "NORMAL": SAMPLE_DIR / "NORMAL-33350-1.jpeg",
}
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

CKA_MATRIX = [
    [1.00, 0.57, 0.40, 0.09, 0.08, 0.07],
    [0.57, 1.00, 0.50, 0.06, 0.05, 0.05],
    [0.40, 0.50, 1.00, 0.47, 0.44, 0.44],
    [0.09, 0.06, 0.47, 1.00, 0.99, 0.99],
    [0.08, 0.05, 0.44, 0.99, 1.00, 1.00],
    [0.07, 0.05, 0.44, 0.99, 1.00, 1.00],
]
CKA_LABELS = ["Patch Embed", "ViT Blk-3", "ViT Blk-7", "ViT Blk-11", "Mamba Blk-1", "Mamba Blk-2"]

_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

# --------------------------------------------------------------------------- #
#  Model loading
# --------------------------------------------------------------------------- #
def build_model() -> ImprovedMedMamba:
    if CKPT_PATH.exists():
        print(f"[*] Loading ImprovedMedMamba checkpoint: {CKPT_PATH.name}")
        model = load_model(str(CKPT_PATH), device=DEVICE)
    else:
        print("[!] Checkpoint not found - using random-init ImprovedMedMamba")
        model = ImprovedMedMamba(num_classes=4)
        model.to(DEVICE).eval()
    print(f"[OK] Model on {DEVICE.upper()} | ViT-B/16 + 2 MedMamba blocks | dim=768")
    return model


# --------------------------------------------------------------------------- #
#  Image utilities
# --------------------------------------------------------------------------- #
def _pil_to_tensor(img: Image.Image) -> torch.Tensor:
    if img.mode != "RGB":
        img = img.convert("RGB")
    return _transform(img).unsqueeze(0).to(DEVICE)


def _pil_to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return "data:image/png;base64," + base64.b64encode(buf.read()).decode()


def _fig_to_b64(fig: plt.Figure) -> str:
    buf = io.BytesIO()
    fig.savefig(buf, format="png", bbox_inches="tight",
                facecolor=fig.get_facecolor(), dpi=100)
    plt.close(fig)
    buf.seek(0)
    return "data:image/png;base64," + base64.b64encode(buf.read()).decode()


def _feat_to_map(feat: torch.Tensor) -> np.ndarray:
    arr  = feat[0].mean(-1).cpu().numpy() if feat.dim() == 3 else feat.mean(-1).cpu().numpy()
    side = int(round(arr.shape[0] ** 0.5))
    arr  = arr[:side * side].reshape(side, side)
    mn, mx = arr.min(), arr.max()
    return ((arr - mn) / (mx - mn + 1e-8)).astype(np.float32)


def _overlay_heatmap(img: Image.Image, hm: np.ndarray,
                     alpha: float = 0.6, cmap: str = "jet") -> Image.Image:
    orig = np.array(img.resize((224, 224)).convert("RGB"), dtype=np.float32) / 255.0
    if hm.shape != (224, 224):
        hm = scipy.ndimage.zoom(hm.astype(np.float32),
                                (224 / hm.shape[0], 224 / hm.shape[1]), order=1)
    hm    = np.clip(hm, 0, 1)
    color = cm.get_cmap(cmap)(hm)[..., :3]
    blend = np.clip((1 - alpha) * orig + alpha * color, 0, 1)
    return Image.fromarray((blend * 255).astype(np.uint8))


def _vit_embed(model: ImprovedMedMamba, tensor: torch.Tensor) -> torch.Tensor:
    """Run ViT patch embedding + CLS + pos embed."""
    B = tensor.shape[0]
    x = model.vit.patch_embed(tensor)
    cls = model.vit.cls_token.expand(B, -1, -1)
    x = torch.cat([cls, x], dim=1)
    return model.vit.pos_drop(x + model.vit.pos_embed)


# --------------------------------------------------------------------------- #
#  Real GradCAM - hooks last ViT block
# --------------------------------------------------------------------------- #
def real_gradcam(model: ImprovedMedMamba, tensor: torch.Tensor,
                 target_cls: Optional[int] = None) -> np.ndarray:
    model.eval()
    acts: List[torch.Tensor] = []
    grads: List[torch.Tensor] = []

    last_vit = model.vit.blocks[-1]

    def fwd_hook(m, inp, out):
        acts.append(out)

    def bwd_hook(m, gin, gout):
        grads.append(gout[0])

    h1 = last_vit.register_forward_hook(fwd_hook)
    h2 = last_vit.register_full_backward_hook(bwd_hook)

    with torch.enable_grad():
        inp = tensor.clone().detach().requires_grad_(True)
        logits = model(inp)
        if target_cls is None:
            target_cls = int(logits.argmax(-1))
        logits[0, target_cls].backward()

    h1.remove(); h2.remove()

    if not acts or not grads:
        return np.zeros((224, 224), dtype=np.float32)

    act  = acts[0][0, 1:].detach()    # (196, D) - drop CLS
    grad = grads[0][0, 1:].detach()   # (196, D)
    weights = grad.mean(0)
    cam = (act.cpu() * weights.cpu()).numpy().sum(axis=-1)   # (196,)
    cam = np.maximum(cam, 0)

    side = int(round(cam.shape[0] ** 0.5))
    cam  = cam[:side * side].reshape(side, side)
    mn, mx = cam.min(), cam.max()
    cam  = (cam - mn) / (mx - mn + 1e-8)
    return scipy.ndimage.zoom(cam.astype(np.float32), 224 / side, order=1)


# --------------------------------------------------------------------------- #
#  Real Attention Rollout
# --------------------------------------------------------------------------- #
def real_attention_rollout(model: ImprovedMedMamba,
                            tensor: torch.Tensor) -> np.ndarray:
    model.eval()
    with torch.no_grad():
        model(tensor)
    attn_maps = model.get_attention_maps()
    if not attn_maps:
        return np.zeros((224, 224), dtype=np.float32)

    N   = attn_maps[0].shape[-1]
    eye = torch.eye(N, device=DEVICE)
    rollout = eye.clone()
    for a in attn_maps:
        am  = a[0].mean(0)
        am  = am + eye
        am  = am / am.sum(-1, keepdim=True)
        rollout = am @ rollout

    cls_attn = rollout[0, 1:].cpu().numpy()
    side     = int(round(cls_attn.shape[0] ** 0.5))
    cls_attn = cls_attn[:side * side].reshape(side, side)
    mn, mx   = cls_attn.min(), cls_attn.max()
    cls_attn = (cls_attn - mn) / (mx - mn + 1e-8)
    return scipy.ndimage.zoom(cls_attn.astype(np.float32), 224 / side, order=1)


# --------------------------------------------------------------------------- #
#  Occlusion Sensitivity
# --------------------------------------------------------------------------- #
def real_occlusion_sensitivity(model: ImprovedMedMamba, img: Image.Image,
                                target_cls: Optional[int], grid: int = 8) -> np.ndarray:
    img_arr = np.array(img.resize((224, 224)).convert("RGB"), dtype=np.uint8)
    step     = 224 // grid
    score_map = np.zeros((grid, grid), dtype=np.float32)
    mean_rgb  = [int(m * 255) for m in IMAGENET_MEAN]

    base_tensor = _pil_to_tensor(img)
    with torch.no_grad():
        base_logits = model(base_tensor)
        if target_cls is None:
            target_cls = int(base_logits.argmax(-1))
        base_score = float(F.softmax(base_logits, dim=-1)[0, target_cls])

    for gy in range(grid):
        for gx in range(grid):
            occ = img_arr.copy()
            y0, y1 = gy * step, min((gy + 1) * step, 224)
            x0, x1 = gx * step, min((gx + 1) * step, 224)
            occ[y0:y1, x0:x1] = mean_rgb
            occ_img = Image.fromarray(occ)
            with torch.no_grad():
                t = _pil_to_tensor(occ_img)
                s = float(F.softmax(model(t), dim=-1)[0, target_cls])
            score_map[gy, gx] = base_score - s

    mn, mx = score_map.min(), score_map.max()
    norm = (score_map - mn) / (mx - mn + 1e-8)
    return scipy.ndimage.zoom(norm.astype(np.float32), 224 / grid, order=1)


# --------------------------------------------------------------------------- #
#  Integrated Gradients
# --------------------------------------------------------------------------- #
def real_integrated_gradients(model: ImprovedMedMamba, tensor: torch.Tensor,
                               target_cls: int, steps: int = 20) -> np.ndarray:
    baseline   = torch.zeros_like(tensor)
    integrated = torch.zeros_like(tensor)

    for k in range(steps):
        alpha = (k + 1) / steps
        inp   = (baseline + alpha * (tensor - baseline)).requires_grad_(True)
        with torch.enable_grad():
            score = model(inp)[0, target_cls]
            score.backward()
        if inp.grad is not None:
            integrated += inp.grad.detach()

    ig = ((tensor - baseline) * integrated / steps)[0]
    ig_map = ig.abs().mean(0).cpu().numpy()
    mn, mx  = ig_map.min(), ig_map.max()
    return ((ig_map - mn) / (mx - mn + 1e-8)).astype(np.float32)


# --------------------------------------------------------------------------- #
#  Neural Journey Figure
# --------------------------------------------------------------------------- #
def make_journey_figure(img_rgb: Image.Image, stage_feats: dict,
                         pred_cls: str) -> str:
    BG = "#0D1117"
    oct_arr = np.array(img_rgb.resize((224, 224)).convert("RGB"))

    stage_keys   = ["patch_embed", "vit_0", "vit_3", "vit_7", "vit_11",
                    "mamba_0", "mamba_1"]
    stage_labels = ["Patch\nEmbed", "ViT\nBlk-0", "ViT\nBlk-3",
                    "ViT\nBlk-7", "ViT\nBlk-11",
                    "MedMamba\nBlk-0", "MedMamba\nBlk-1"]

    feat_maps = {}
    for k in stage_keys:
        if k in stage_feats:
            feat_maps[k] = _feat_to_map(stage_feats[k])

    n_cols = len(stage_keys) + 1
    fig, axes = plt.subplots(2, n_cols, figsize=(20, 7), dpi=95)
    fig.patch.set_facecolor(BG)
    fig.suptitle(f"RetViM Neural Journey - {pred_cls}  [Real Model * ImprovedMedMamba * dim=768]",
                 color="white", fontsize=11, fontweight="bold", y=1.01)

    for col in range(n_cols):
        for row in range(2):
            ax = axes[row, col]
            ax.set_facecolor("#161B22")
            ax.axis("off")
            if row == 0:
                ax.set_title("Input" if col == 0 else stage_labels[col - 1],
                             color="white", fontsize=6.5, pad=2)

    axes[0, 0].imshow(oct_arr)
    for ci, key in enumerate(stage_keys):
        if key in feat_maps:
            cmap_n = "plasma" if key.startswith("vit") or key == "patch_embed" else "inferno"
            axes[0, ci + 1].imshow(feat_maps[key], cmap=cmap_n, interpolation="bilinear")

    axes[1, 0].imshow(oct_arr)
    for ci, key in enumerate(stage_keys):
        if key in feat_maps:
            fm = feat_maps[key]
            zoom_factor = 224 / fm.shape[0]
            fm_up = scipy.ndimage.zoom(fm, zoom_factor, order=1)
            ov = _overlay_heatmap(img_rgb, fm_up, alpha=0.55)
            axes[1, ci + 1].imshow(np.array(ov))

    plt.tight_layout(pad=0.6)
    b64 = _fig_to_b64(fig)
    gc.collect()
    return b64


# --------------------------------------------------------------------------- #
#  Deep Analysis - Attention Heads + Mamba Internals
# --------------------------------------------------------------------------- #
_CACHE: dict = {}
_CLASS_PROTOTYPES: Optional[dict] = None


def _extract_deep_analysis(model: ImprovedMedMamba, tensor: torch.Tensor) -> dict:
    global _CLASS_PROTOTYPES

    num_vit   = len(model.vit.blocks)          # 12
    num_mamba = len(model.medmamba_blocks)      # 2
    num_heads = model.vit.blocks[0].attn.num_heads  # 12

    cls_tokens    = []
    magnitudes    = []
    frozen_acts   = []
    trainable_acts = []

    with torch.no_grad():
        x = _vit_embed(model, tensor)

        for i, blk in enumerate(model.vit.blocks):
            x = blk(x)
            cls_tokens.append(x[0, 0].clone())
            magnitudes.append(torch.norm(x[0, 1:], dim=-1).mean().item())
            act_vals = x[0, 1:].flatten().cpu().numpy()
            if i < 6:
                frozen_acts.append(act_vals)
            else:
                trainable_acts.append(act_vals)

        x = model.vit.norm(x)

        # Mamba blocks - enable internal storage
        model.enable_mamba_store(True)
        for mi, blk in enumerate(model.medmamba_blocks):
            x = blk(x)
            cls_tokens.append(x[0, 0].clone())
            magnitudes.append(torch.norm(x[0, 1:], dim=-1).mean().item())
        model.enable_mamba_store(False)

    # --- Attention heads ---
    attn_maps = model.get_attention_maps()
    attention_heads = []
    head_entropy    = []
    layer_entropy   = []

    for li, attn in enumerate(attn_maps):
        layer_heads    = []
        layer_head_ent = []
        layer_ent_sum  = 0.0
        for hi in range(num_heads):
            cls_attn = attn[0, hi, 0, 1:].cpu()   # (196,)
            head_map = cls_attn.reshape(14, 14)
            hmn, hmx = head_map.min(), head_map.max()
            head_map = ((head_map - hmn) / (hmx - hmn + 1e-8)).numpy().tolist()
            layer_heads.append(head_map)
            p   = cls_attn + 1e-10; p /= p.sum()
            ent = -(p * torch.log2(p)).sum().item()
            layer_head_ent.append(round(ent, 4))
            layer_ent_sum += ent
        attention_heads.append(layer_heads)
        head_entropy.append(layer_head_ent)
        layer_entropy.append(round(layer_ent_sum / num_heads, 4))

    # --- Mamba internals ---
    mamba_internals = model.get_mamba_internals()
    # Ensure conv_map / ssm_map are 2-D (list of lists)
    for blk_data in mamba_internals:
        for key in ("conv_map", "ssm_map", "fusion_map", "conv_ssm_ratio"):
            raw = blk_data.get(key, [])
            # If 1-D flat list, try to square-reshape to 14x14
            if raw and not isinstance(raw[0], list):
                side = int(round(len(raw) ** 0.5))
                if side * side == len(raw):
                    blk_data[key] = [raw[y * side:(y + 1) * side] for y in range(side)]
                else:
                    blk_data[key] = []

    # --- CLS class-similarity trajectory ---
    if _CLASS_PROTOTYPES is None:
        _CLASS_PROTOTYPES = {}
        with torch.no_grad():
            for ci, cn in enumerate(CLASS_NAMES):
                syn_img = draw_synthetic_oct(cn)
                syn_t   = _pil_to_tensor(syn_img)
                sx      = _vit_embed(model, syn_t)
                for blk in model.vit.blocks:
                    sx = blk(sx)
                sx = model.vit.norm(sx)
                for blk in model.medmamba_blocks:
                    sx = blk(sx)
                _CLASS_PROTOTYPES[cn] = sx[0, 0].clone()

    cls_similarity: dict = {}
    for cn, proto in _CLASS_PROTOTYPES.items():
        sims = [round(F.cosine_similarity(ct.unsqueeze(0), proto.unsqueeze(0)).item(), 4)
                for ct in cls_tokens]
        cls_similarity[cn] = sims

    # --- Activation histograms ---
    f_arr = np.concatenate(frozen_acts)   if frozen_acts   else np.zeros(100)
    t_arr = np.concatenate(trainable_acts) if trainable_acts else np.zeros(100)
    f_hist, f_bins = np.histogram(f_arr, bins=50)
    t_hist, _      = np.histogram(t_arr, bins=f_bins)
    f_hist = (f_hist / (f_hist.max() + 1e-8)).tolist()
    t_hist = (t_hist / (t_hist.max() + 1e-8)).tolist()

    return {
        "attention_heads":  attention_heads,
        "head_entropy":     head_entropy,
        "layer_entropy":    layer_entropy,
        "layer_magnitude":  magnitudes,
        "cls_similarity":   cls_similarity,
        "frozen_hist":      f_hist,
        "trainable_hist":   t_hist,
        "hist_bins":        f_bins.tolist(),
        "mamba_internals":  mamba_internals,
        "num_heads":        num_heads,
        "num_mamba_blocks": num_mamba,
    }


# --------------------------------------------------------------------------- #
#  Synthetic OCT generator (used for /sample endpoint)
# --------------------------------------------------------------------------- #
def draw_synthetic_oct(class_name: str, size: int = 224) -> Image.Image:
    cn  = class_name.upper()
    rng = np.random.RandomState(CLASS_NAMES.index(cn) * 7 + 13)
    arr = np.zeros((size, size), dtype=np.int16)
    rpe_y   = int(size * 0.68)
    ilm_y   = int(size * 0.28)
    inner_y = int(size * 0.38)

    for y in range(size):
        arr[y, :] = int(10 + (y / size) * 15)

    for x in range(size):
        wave = int(4 * math.sin(x * 0.05))
        arr[max(0, ilm_y + wave - 2): ilm_y + wave + 2, x]   = 200
        arr[max(0, inner_y + wave - 3): inner_y + wave + 3, x] = 140
        arr[max(0, rpe_y - wave - 3): rpe_y - wave + 3, x]    = 220
        clen = min(15, size - (rpe_y - wave + 3))
        if clen > 0:
            arr[rpe_y - wave + 3: rpe_y - wave + 3 + clen, x] = \
                np.clip(rng.randint(40, 80, clen), 30, 90)

    arr[inner_y:rpe_y, :] = np.clip(
        arr[inner_y:rpe_y, :] + rng.randint(0, 20, (rpe_y - inner_y, size), dtype=np.int16),
        0, 255)

    if cn == "CNV":
        cx = size // 2
        for dx in range(-40, 41):
            for dy in range(-12, 13):
                if dx * dx / 1600 + dy * dy / 144 <= 1:
                    px, py = cx + dx, rpe_y - 10 + dy
                    if 0 <= px < size and 0 <= py < size:
                        arr[py, px] = int(np.clip(arr[py, px] + 100, 0, 255))
    elif cn == "DME":
        for cx, cy, rw, rh in [(size//2-20, inner_y+15, 18, 10),
                                (size//2+25, inner_y+25, 14,  8),
                                (size//2,    inner_y+35, 20, 12)]:
            for dx in range(-rw, rw + 1):
                for dy in range(-rh, rh + 1):
                    if dx * dx / (rw * rw) + dy * dy / (rh * rh) <= 1:
                        px, py = cx + dx, cy + dy
                        if 0 <= px < size and 0 <= py < size:
                            arr[py, px] = 8
    elif cn == "DRUSEN":
        rng2 = np.random.RandomState(42)
        for _ in range(12):
            cx = rng2.randint(30, size - 30); w = rng2.randint(8, 20); h2 = rng2.randint(4, 10)
            for dx in range(-w, w + 1):
                for dy in range(-h2, 0):
                    if dx * dx / (w * w) + dy * dy / (h2 * h2) <= 1:
                        px, py = cx + dx, rpe_y + dy - 2
                        if 0 <= px < size and 0 <= py < size:
                            arr[py, px] = int(np.clip(arr[py, px] + 80, 0, 255))
    elif cn == "NORMAL":
        arr = scipy.ndimage.gaussian_filter(arr.astype(float), sigma=1.2).astype(np.int16)

    noise = rng.randint(-12, 12, arr.shape, dtype=np.int16)
    arr   = np.clip(arr + noise, 0, 255).astype(np.uint8)
    return Image.fromarray(arr, "L").convert("RGB")


# --------------------------------------------------------------------------- #
#  Full analysis pipeline
# --------------------------------------------------------------------------- #
def run_analysis(model: ImprovedMedMamba, img: Image.Image) -> dict:
    """Run complete RetViM analysis. ALL results are real PyTorch computations."""
    img_rgb = img.convert("RGB")
    img_key = hashlib.md5(img_rgb.resize((64, 64)).tobytes()).hexdigest()
    if img_key in _CACHE and "deep_analysis" in _CACHE[img_key]:
        return _CACHE[img_key]

    tensor = _pil_to_tensor(img_rgb)

    # 1. Real classification forward pass
    t0 = time.perf_counter()
    with torch.no_grad():
        logits = model(tensor)
    latency_ms = round((time.perf_counter() - t0) * 1000, 2)

    probs_t    = F.softmax(logits, dim=-1)[0].cpu().tolist()
    pred_idx   = int(np.argmax(probs_t))
    pred_cls   = CLASS_NAMES[pred_idx]
    confidence = probs_t[pred_idx]

    # 2. Real GradCAM
    print(f"  [*] GradCAM backprop -> {pred_cls} ...")
    gcam_map  = real_gradcam(model, tensor, target_cls=pred_idx)
    gcam_pil  = Image.fromarray(
        (cm.get_cmap("jet")(gcam_map)[..., :3] * 255).astype(np.uint8))
    gcam_b64  = _pil_to_b64(gcam_pil)

    # 3. Attention rollout
    print("  [*] Attention rollout ...")
    rollout_map = real_attention_rollout(model, tensor)
    rollout_pil = Image.fromarray(
        (cm.get_cmap("viridis")(rollout_map)[..., :3] * 255).astype(np.uint8))
    rollout_b64 = _pil_to_b64(rollout_pil)

    # 4. Overlays
    overlay_b64         = _pil_to_b64(_overlay_heatmap(img_rgb, gcam_map,    alpha=0.6))
    rollout_overlay_b64 = _pil_to_b64(_overlay_heatmap(img_rgb, rollout_map, alpha=0.55, cmap="viridis"))

    # 5. Occlusion sensitivity (6x6 on CPU to keep latency reasonable)
    _occ_grid = 6 if DEVICE == "cpu" else 8
    print(f"  [*] Occlusion sensitivity ({_occ_grid}x{_occ_grid}) ...")
    occ_map  = real_occlusion_sensitivity(model, img_rgb, target_cls=pred_idx, grid=_occ_grid)
    occ_pil  = Image.fromarray(
        (cm.get_cmap("RdYlGn")(occ_map)[..., :3] * 255).astype(np.uint8))
    occ_b64  = _pil_to_b64(occ_pil)

    # 6. Integrated Gradients (15 steps on CPU, 25 on GPU)
    _ig_steps = 15 if DEVICE == "cpu" else 25
    print(f"  [*] Integrated Gradients ({_ig_steps} steps) ...")
    ig_map  = real_integrated_gradients(model, tensor, target_cls=pred_idx, steps=_ig_steps)
    ig_pil  = Image.fromarray(
        (cm.get_cmap("inferno")(ig_map)[..., :3] * 255).astype(np.uint8))
    ig_b64  = _pil_to_b64(ig_pil)

    # 7. Intermediate feature maps
    print("  [*] Extracting feature maps ...")
    with torch.no_grad():
        stage_feats = model.get_intermediate_features(tensor)

    # 8. Neural journey figure
    print("  [*] Neural journey figure ...")
    journey_b64 = make_journey_figure(img_rgb, stage_feats, pred_cls)

    # 9. Deep analysis (attention heads + Mamba internals)
    print("  [*] Deep analysis (heads + Mamba internals) ...")
    deep = _extract_deep_analysis(model, tensor)

    prob_dict = {CLASS_NAMES[i]: round(probs_t[i], 6) for i in range(4)}

    result = {
        "prediction":           pred_cls,
        "confidence":           round(confidence, 6),
        "probabilities":        probs_t,
        "prob_dict":            prob_dict,
        "gcam_b64":             gcam_b64,
        "rollout_b64":          rollout_b64,
        "overlay_b64":          overlay_b64,
        "rollout_overlay_b64":  rollout_overlay_b64,
        "occ_b64":              occ_b64,
        "ig_b64":               ig_b64,
        "neural_journey_b64":   journey_b64,
        "cka_matrix":           CKA_MATRIX,
        "latency_ms":           latency_ms,
        "model":                "RetViM - ImprovedMedMamba (ViT-B/16 + 2 SS-Conv-SSM)",
        "device":               DEVICE,
        "trained_weights":      CKPT_PATH.exists(),
        "deep_analysis":        deep,
    }
    _CACHE[img_key] = result
    print(f"  [OK] {pred_cls} ({confidence*100:.2f}%) in {latency_ms}ms")
    return result


# --------------------------------------------------------------------------- #
#  Build model at startup
# --------------------------------------------------------------------------- #
print("[*] Building RetViM model ...")
MODEL = build_model()
param_count   = sum(p.numel() for p in MODEL.parameters())
HAS_WEIGHTS   = CKPT_PATH.exists()
print(f"[OK] Ready on {DEVICE.upper()} | {param_count:,} parameters")

# --------------------------------------------------------------------------- #
#  FastAPI application
# --------------------------------------------------------------------------- #
app = FastAPI(
    title="RetViM Real Inference Server",
    description="Real PyTorch inference - ImprovedMedMamba (ViT-B/16 + 2 SS-Conv-SSM blocks)",
    version="4.0.0",
)
app.add_middleware(
    CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"],
)

if (BASE_DIR / "images").exists():
    app.mount("/images", StaticFiles(directory="images"), name="images")
if (BASE_DIR / "public").exists():
    app.mount("/public", StaticFiles(directory="public"), name="public")


@app.get("/", response_class=HTMLResponse)
async def root():
    if HTML_FILE.exists():
        html = HTML_FILE.read_text(encoding="utf-8")
        # Replace any HF Space URL with local address (robust regex)
        html = re.sub(
            r"const API_BASE\s*=\s*'https?://[^']*';",
            f"const API_BASE = 'http://{HOST}:{PORT}';",
            html,
        )
        return HTMLResponse(content=html)
    return HTMLResponse("<h1>retvim-final.html not found.</h1>", status_code=404)


@app.get("/health")
async def health():
    return {
        "status":          "ok",
        "model_loaded":    True,
        "trained_weights": HAS_WEIGHTS,
        "device":          DEVICE,
        "model_params":    param_count,
        "uptime_s":        round(time.time() - START_TIME, 1),
        "inference":       "real PyTorch - ImprovedMedMamba",
        "architecture":    "ViT-B/16 (dim=768, 12 blocks, 12 heads) + 2 SS-Conv-SSM blocks",
        "val_acc":         "96.68%" if HAS_WEIGHTS else "random init",
    }


@app.get("/classes")
async def list_classes():
    return {
        "classes": CLASS_NAMES,
        "descriptions": {
            "CNV":    "Choroidal Neovascularization - abnormal blood vessel growth below RPE",
            "DME":    "Diabetic Macular Edema - cystoid intraretinal fluid accumulation",
            "DRUSEN": "Drusen deposits - sub-RPE yellow deposits (dry AMD indicator)",
            "NORMAL": "Healthy retina - all layers intact, no pathological findings",
        },
    }


@app.get("/cka")
async def cka():
    return {"matrix": CKA_MATRIX, "labels": CKA_LABELS}


@app.get("/sample-image/{class_name}")
async def sample_image(class_name: str):
    """Serve the raw real OCT sample JPEG for thumbnail display."""
    cn = class_name.upper()
    path = SAMPLE_IMAGES.get(cn)
    if not path or not path.exists():
        return JSONResponse(status_code=404, content={"detail": "Image not found"})
    return FileResponse(str(path), media_type="image/jpeg",
                        headers={"Cache-Control": "public, max-age=86400"})


@app.get("/sample/{class_name}")
async def sample(class_name: str):
    """Real model inference on a real OCT sample image."""
    cn = class_name.upper()
    if cn not in CLASS_NAMES:
        return JSONResponse(status_code=404,
                            content={"detail": f"Unknown class. Valid: {CLASS_NAMES}"})
    # Use real sample image when available, else fall back to synthetic
    path = SAMPLE_IMAGES.get(cn)
    if path and path.exists():
        img = Image.open(path).convert("RGB")
    else:
        img = draw_synthetic_oct(cn)
    return run_analysis(MODEL, img)


@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    mode: Optional[str] = Form(None),
    mode_q: Optional[str] = Query(None, alias="mode"),
):
    """Upload a real OCT image -> get real model inference.
    mode = 'doctor' | 'researcher'  (send as form field or ?mode= query param)
    """
    # Accept mode from either form field or query param
    effective_mode = (mode or mode_q or "doctor").lower()

    raw = await file.read()
    if not raw:
        return JSONResponse(status_code=400, content={"detail": "Empty file."})
    try:
        img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception as e:
        return JSONResponse(status_code=400, content={"detail": f"Invalid image: {e}"})

    data = run_analysis(MODEL, img)

    if effective_mode == "doctor":
        return {
            "prediction":    data["prediction"],
            "confidence":    data["confidence"],
            "probabilities": data["probabilities"],
            "prob_dict":     data["prob_dict"],
            "gcam_b64":      data["gcam_b64"],
            "overlay_b64":   data["overlay_b64"],
            "latency_ms":    data["latency_ms"],
        }
    # researcher mode: return everything
    return data


@app.get("/figures/{fig_type}")
async def figures(fig_type: str, class_name: str = Query("CNV")):
    cn = class_name.upper()
    if cn not in CLASS_NAMES:
        return JSONResponse(status_code=400, content={"detail": "Invalid class."})
    data   = run_analysis(MODEL, draw_synthetic_oct(cn))
    b64_key = {
        "neural_journey": "neural_journey_b64",
        "gradcam":        "gcam_b64",
        "overlay":        "overlay_b64",
        "rollout":        "rollout_b64",
    }.get(fig_type)
    if not b64_key:
        return JSONResponse(status_code=404, content={"detail": "Unknown fig_type."})
    raw_b64 = data[b64_key].split(",", 1)[-1]
    return Response(content=base64.b64decode(raw_b64), media_type="image/png")


# --------------------------------------------------------------------------- #
#  Entry point
# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    print()
    print("=" * 70)
    print("  RetViM - Real Inference Server v4.0")
    print("=" * 70)
    print(f"  Frontend  : http://{HOST}:{PORT}/")
    print(f"  API docs  : http://{HOST}:{PORT}/docs")
    print(f"  Health    : http://{HOST}:{PORT}/health")
    print(f"  Device    : {DEVICE.upper()}")
    print(f"  Model     : ImprovedMedMamba - ViT-B/16 + 2 SS-Conv-SSM - dim=768")
    print(f"  Weights   : {'[OK] Real checkpoint (96.68% val_acc)' if HAS_WEIGHTS else '[!] No checkpoint - random init'}")
    print(f"  Params    : {param_count:,}")
    print("=" * 70)
    print()
    uvicorn.run("server:app", host=HOST, port=PORT, reload=False)