inference.py

"""
inference.py – RetViMVisualizer
All research-figure generation for the RetViM paper.
"""

import io
import base64
import hashlib
import json
import time
import os
from typing import Optional

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import LinearSegmentedColormap

from torchvision import transforms
from sklearn.decomposition import PCA
import scipy.ndimage

from model import load_model

# ── constants ──────────────────────────────────────────────────────────────────

CLASS_NAMES = ["CNV", "DME", "DRUSEN", "NORMAL"]

# Real checkpoint path
REAL_CKPT = "./improved-medmamba-epoch=19-val_acc=0.9668.ckpt"
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]
BG_COLOR = "#0D1117"
CACHE_TTL = 600  # 10 minutes

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


# ── synthetic OCT generator ────────────────────────────────────────────────────

def draw_synthetic_oct(class_name: str, size: int = 224) -> Image.Image:
    """
    Draw a realistic grayscale OCT image with class-specific pathology.
    Returns RGB PIL image (for compatibility with transforms).
    """
    img = Image.new("L", (size, size), color=10)
    draw = ImageDraw.Draw(img)

    # Retinal layers (bottom = RPE, top = ILM)
    rpe_y   = int(size * 0.68)
    ilm_y   = int(size * 0.28)
    inner_y = int(size * 0.38)

    # Background gradient (vitreous) – use int16 to avoid overflow
    arr = np.zeros((size, size), dtype=np.int16)
    for y in range(size):
        val = int(10 + (y / size) * 15)
        arr[y, :] = val

    # Draw retinal layers
    for x in range(size):
        wave = int(4 * np.sin(x * 0.05))
        # ILM
        arr[max(0, ilm_y + wave - 2): ilm_y + wave + 2, x] = 200
        # Inner nuclear / plexiform
        arr[max(0, inner_y + wave - 3): inner_y + wave + 3, x] = 140
        # RPE line
        arr[max(0, rpe_y - wave - 3): rpe_y - wave + 3, x] = 220
        # Choroid
        choroid_len = min(15, size - (rpe_y - wave + 3))
        if choroid_len > 0:
            arr[rpe_y - wave + 3: rpe_y - wave + 3 + choroid_len, x] = \
                np.clip(np.random.randint(40, 80, choroid_len), 30, 90)

    # Subretinal space noise
    arr[inner_y:rpe_y, :] = np.clip(
        arr[inner_y:rpe_y, :] + np.random.randint(0, 20, (rpe_y - inner_y, size), dtype=np.int16),
        0, 255
    )

    cn = class_name.upper()
    if cn == "CNV":
        # Subretinal membrane / neovascular complex
        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))
        # Subretinal fluid
        for dx in range(-55, 56):
            for dy in range(-8, 9):
                px, py = cx + dx, inner_y + 20 + dy
                if 0 <= px < size and 0 <= py < size:
                    arr[py, px] = int(np.clip(arr[py, px] - 30, 0, 255))

    elif cn == "DME":
        # Cystoid spaces (intraretinal fluid)
        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  # dark cysts

    elif cn == "DRUSEN":
        # Drusen deposits along RPE
        np.random.seed(42)
        for _ in range(12):
            cx = np.random.randint(30, size - 30)
            w  = np.random.randint(8, 20)
            h  = np.random.randint(4, 10)
            for dx in range(-w, w + 1):
                for dy in range(-h, 0):
                    if dx * dx / (w * w) + dy * dy / (h * h) <= 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":
        # Clean retina; slight smoothing
        from scipy.ndimage import gaussian_filter
        arr = gaussian_filter(arr.astype(float), sigma=1.2).astype(np.uint8)

    # Add mild speckle noise (OCT characteristic) – arr is int16, safe to add
    noise = np.random.randint(-12, 12, arr.shape, dtype=np.int16)
    arr = np.clip(arr + noise, 0, 255).astype(np.uint8)

    gray = Image.fromarray(arr, mode="L")
    return gray.convert("RGB")


# ── helpers ────────────────────────────────────────────────────────────────────

def _to_tensor(image: Image.Image, device: str) -> torch.Tensor:
    if image.mode != "RGB":
        image = image.convert("RGB")
    return _transform(image).unsqueeze(0).to(device)


def _fig_to_b64(fig: plt.Figure) -> str:
    buf = io.BytesIO()
    fig.savefig(buf, format="png", facecolor=fig.get_facecolor(),
                bbox_inches="tight", dpi=fig.get_dpi())
    plt.close(fig)
    buf.seek(0)
    return base64.b64encode(buf.read()).decode("utf-8")


def _pil_to_b64(image: Image.Image) -> str:
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    buf.seek(0)
    return base64.b64encode(buf.read()).decode("utf-8")


def _b64_to_pil(b64: str) -> Image.Image:
    return Image.open(io.BytesIO(base64.b64decode(b64)))


def _feature_to_map(feat: torch.Tensor, h: int = 14, w: int = 14) -> np.ndarray:
    """(1, N, D) → normalised (H, W) spatial mean activation."""
    arr = feat[0].mean(-1).cpu().numpy()  # (N,)
    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(image: Image.Image, heatmap: np.ndarray,
                     alpha: float = 0.6, cmap_name: str = "jet") -> Image.Image:
    import scipy.ndimage
    orig = np.array(image.resize((224, 224)).convert("RGB"), dtype=np.float32) / 255.0
    # Upsample heatmap to match image size if needed
    if heatmap.shape != (224, 224):
        zoom_y = 224 / heatmap.shape[0]
        zoom_x = 224 / heatmap.shape[1]
        heatmap = scipy.ndimage.zoom(heatmap.astype(np.float32), (zoom_y, zoom_x), order=1)
        heatmap = np.clip(heatmap, 0, 1)
    cmap = cm.get_cmap(cmap_name)(heatmap)[..., :3]
    blended = np.clip((1 - alpha) * orig + alpha * cmap, 0, 1)
    return Image.fromarray((blended * 255).astype(np.uint8))


def _cache_path(key: str) -> str:
    return f"/tmp/retvim_{key}.json"


def _load_cache(key: str) -> Optional[dict]:
    p = _cache_path(key)
    try:
        if os.path.exists(p):
            with open(p) as f:
                data = json.load(f)
            if time.time() - data.get("_ts", 0) < CACHE_TTL:
                return data
    except Exception:
        pass
    return None


def _save_cache(key: str, data: dict):
    data["_ts"] = time.time()
    try:
        with open(_cache_path(key), "w") as f:
            json.dump(data, f)
    except Exception:
        pass


# ── RetViMVisualizer ───────────────────────────────────────────────────────────

class RetViMVisualizer:
    """
    Unified inference + visualisation class for the RetViM paper.
    """

    # ── 1. __init__ ────────────────────────────────────────────────────────────

    def __init__(self, model_path: str, device: str = "cpu"):
        self.device = device
        self.model = load_model(model_path, num_classes=4, device=device)
        self.model.eval()
        self.class_names = CLASS_NAMES
        self.transform = _transform

    # ── 2. predict ─────────────────────────────────────────────────────────────

    def predict(self, image: Image.Image) -> dict:
        """Returns prediction, confidence, probabilities, latency_ms."""
        tensor = _to_tensor(image, self.device)
        t0 = time.perf_counter()
        with torch.no_grad():
            logits = self.model(tensor)
        latency_ms = round((time.perf_counter() - t0) * 1000, 2)

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

        return {
            "prediction":    self.class_names[pred_idx],
            "confidence":    round(probs[pred_idx], 6),
            "probabilities": {c: round(p, 6) for c, p in zip(self.class_names, probs)},
            "latency_ms":    latency_ms,
        }

    # ── 3. generate_neural_journey ─────────────────────────────────────────────

    def generate_neural_journey(self, image: Image.Image, class_name: str) -> str:
        """
        3-row figure:
          Row 1 – Input OCT + feature maps at each stage
          Row 2 – Attention overlay heatmaps (jet, alpha=0.6)
          Row 3 – Activation statistics line chart
        Returns base64 PNG string.
        """
        tensor = _to_tensor(image, self.device)
        with torch.no_grad():
            feats = self.model.get_intermediate_features(tensor)
            # trigger attention storage
            _ = self.model(tensor)

        stage_keys  = ["patch_embed", "vit_0", "vit_3", "vit_7", "vit_11",
                        "mamba_0", "mamba_1"]
        stage_labels = ["Patch\nEmbed", "ViT\nBlock 0", "ViT\nBlock 3",
                        "ViT\nBlock 7", "ViT\nBlock 11",
                        "MedMamba\nBlock 0", "MedMamba\nBlock 1"]
        vit_cmap    = "plasma"
        mamba_cmap  = "inferno"

        maps = {k: _feature_to_map(feats[k]) for k in stage_keys if k in feats}

        n_cols = len(stage_keys) + 1       # +1 for input image
        fig, axes = plt.subplots(3, n_cols, figsize=(20, 8), dpi=120)
        fig.patch.set_facecolor(BG_COLOR)

        oct_rgb = np.array(image.resize((224, 224)).convert("RGB"))

        # -- Row 0: feature maps --------------------------------------------------
        axes[0, 0].imshow(oct_rgb)
        axes[0, 0].set_title("Input OCT", color="white", fontsize=7, pad=3)
        axes[0, 0].axis("off")

        for ci, key in enumerate(stage_keys):
            ax = axes[0, ci + 1]
            if key in maps:
                cmap = vit_cmap if key.startswith("vit") else mamba_cmap
                ax.imshow(maps[key], cmap=cmap, interpolation="bilinear")
            ax.set_title(stage_labels[ci], color="white", fontsize=6, pad=2)
            ax.axis("off")

        # -- Row 1: attention overlays --------------------------------------------
        axes[1, 0].imshow(oct_rgb)
        axes[1, 0].set_title("Input OCT", color="white", fontsize=7, pad=3)
        axes[1, 0].axis("off")

        for ci, key in enumerate(stage_keys):
            ax = axes[1, ci + 1]
            if key in maps:
                overlay = _overlay_heatmap(image, maps[key], alpha=0.6, cmap_name="jet")
                ax.imshow(np.array(overlay))
            ax.set_title(stage_labels[ci], color="white", fontsize=6, pad=2)
            ax.axis("off")

        # -- Row 2: activation statistics line chart ------------------------------
        ax_stat = axes[2, :]
        for ax in ax_stat:
            ax.set_visible(False)
        ax_line = fig.add_axes([0.05, 0.02, 0.90, 0.24])
        ax_line.set_facecolor("#161B22")

        vit_means   = [feats[k][0].mean(-1).mean().item() for k in stage_keys
                       if k.startswith("vit") and k in feats]
        mamba_means = [feats[k][0].mean(-1).mean().item() for k in stage_keys
                       if k.startswith("mamba") and k in feats]

        if vit_means:
            ax_line.plot(range(len(vit_means)), vit_means, "o-",
                         color="cyan", label="ViT stages", linewidth=1.5, markersize=4)
        if mamba_means:
            x0 = len(vit_means)
            ax_line.plot(range(x0, x0 + len(mamba_means)), mamba_means, "s-",
                         color="orange", label="MedMamba stages", linewidth=1.5, markersize=4)

        ax_line.set_xlabel("Stage", color="white", fontsize=8)
        ax_line.set_ylabel("Mean Activation", color="white", fontsize=8)
        ax_line.set_title(f"Activation Statistics – {class_name}", color="white", fontsize=9)
        ax_line.tick_params(colors="white")
        for spine in ax_line.spines.values():
            spine.set_color("#30363D")
        ax_line.legend(facecolor="#161B22", labelcolor="white", fontsize=7)

        fig.suptitle(f"RetViM Neural Journey – {class_name}", color="white",
                     fontsize=12, fontweight="bold", y=0.98)

        b64 = _fig_to_b64(fig)
        return b64

    # ── 4. generate_gradcam ────────────────────────────────────────────────────

    def generate_gradcam_fast(self, image: Image.Image, target_class: Optional[int] = None) -> dict:
        """
        Fast XAI: GradCAM, GradCAM++, Attention Rollout, Overlay only.
        Skips slow methods (RISE, IG, Occlusion). Used for real-time /predict.
        Returns {gcam_b64, gcam2_b64, rollout_b64, overlay_b64}.
        """
        tensor = _to_tensor(image, self.device)

        gcam_map, gcam2_map = self._grad_cam_both(tensor, target_class)
        rollout_map = self._attention_rollout(tensor)
        overlay_img = _overlay_heatmap(image, gcam_map, alpha=0.6)

        def _map_to_b64(m: np.ndarray, cmap_name: str = "jet") -> str:
            fig, ax = plt.subplots(figsize=(4, 4))
            fig.patch.set_facecolor(BG_COLOR)
            ax.imshow(m, cmap=cmap_name)
            ax.axis("off")
            return _fig_to_b64(fig)

        return {
            "gcam_b64":    _map_to_b64(gcam_map,    "jet"),
            "gcam2_b64":   _map_to_b64(gcam2_map,   "jet"),
            "rollout_b64": _map_to_b64(rollout_map, "viridis"),
            "overlay_b64": _pil_to_b64(overlay_img),
        }

    def generate_gradcam(self, image: Image.Image, target_class: Optional[int] = None) -> dict:
        """
        Returns {gcam_b64, gcam2_b64, rollout_b64, occ_b64, ig_b64, rise_b64, overlay_b64}.
        Uses multiple XAI methods for the explainability suite.
        """
        tensor = _to_tensor(image, self.device)

        # -- GradCAM & GradCAM++ ----------------------------------------------
        gcam_map, gcam2_map = self._grad_cam_both(tensor, target_class)

        # -- Attention Rollout -------------------------------------------------
        rollout_map = self._attention_rollout(tensor)

        # -- Occlusion Sensitivity (8×8 grid) ----------------------------------
        occ_map = self._occlusion_sensitivity(image, target_class, grid=8)

        # -- Integrated Gradients ----------------------------------------------
        ig_map = self._integrated_gradients(tensor, target_class, steps=20)

        # -- RISE --------------------------------------------------------------
        rise_map = self._rise(tensor, target_class, N=200, s=8, p1=0.5)

        # -- Overlay on input -------------------------------------------------
        overlay_img = _overlay_heatmap(image, gcam_map, alpha=0.6)

        def _map_to_b64(m: np.ndarray, cmap_name: str = "jet") -> str:
            fig, ax = plt.subplots(figsize=(4, 4))
            fig.patch.set_facecolor(BG_COLOR)
            ax.imshow(m, cmap=cmap_name)
            ax.axis("off")
            return _fig_to_b64(fig)

        return {
            "gcam_b64":    _map_to_b64(gcam_map,    "jet"),
            "gcam2_b64":   _map_to_b64(gcam2_map,   "jet"),
            "rollout_b64": _map_to_b64(rollout_map, "viridis"),
            "occ_b64":     _map_to_b64(occ_map,     "RdYlGn"),
            "ig_b64":      _map_to_b64(ig_map,      "plasma"),
            "rise_b64":    _map_to_b64(rise_map,    "magma"),
            "overlay_b64": _pil_to_b64(overlay_img),
        }

    def _grad_cam_both(self, tensor: torch.Tensor, target_class: Optional[int]) -> tuple[np.ndarray, np.ndarray]:
        """Gradient-weighted spatial map via last ViT block, returns (GradCAM, GradCAM++)."""
        tensor = tensor.clone().requires_grad_(True)
        # Hook activations & grads from last ViT block output
        acts, grads = [], []

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

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

        # Real model has vit.blocks (ViT-Base/16 inside self.vit)
        last_vit = self.model.vit.blocks[-1]
        h_fwd = last_vit.register_forward_hook(fwd_hook)
        h_bwd = last_vit.register_full_backward_hook(bwd_hook)

        logits = self.model(tensor)
        if target_class is None:
            target_class = int(logits.argmax(-1))
        logits[0, target_class].backward()

        h_fwd.remove()
        h_bwd.remove()

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

        act  = acts[0][0, 1:]   # (N, D) – drop CLS
        grad = grads[0][0, 1:]  # (N, D)
        
        # Standard GradCAM
        weights1 = grad.mean(0)  # (D,)
        cam1 = (act * weights1).sum(-1).cpu().numpy()  # (N,)

        # GradCAM++
        grad_2 = grad.pow(2)
        grad_3 = grad.pow(3)
        alpha = grad_2 / (2 * grad_2 + (act * grad_3).sum(dim=0, keepdim=True) + 1e-8)
        weights2 = (alpha * torch.relu(grad)).sum(0)
        cam2 = (act * weights2).sum(-1).cpu().numpy()

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

        return _process_cam(cam1), _process_cam(cam2)

    def _integrated_gradients(self, tensor: torch.Tensor, target_class: Optional[int], steps: int = 20) -> np.ndarray:
        """Integrated gradients computation."""
        with torch.no_grad():
            logits = self.model(tensor)
        if target_class is None:
            target_class = int(logits.argmax(-1))
            
        baseline = torch.zeros_like(tensor)
        scaled_inputs = [baseline + (float(i) / steps) * (tensor - baseline) for i in range(0, steps + 1)]
        
        grads = []
        for inp in scaled_inputs:
            inp = inp.clone().requires_grad_(True)
            out = self.model(inp)
            out[0, target_class].backward()
            grads.append(inp.grad.detach())
            
        avg_grads = torch.stack(grads).mean(dim=0)
        integrated_grad = (tensor - baseline) * avg_grads
        
        # Convert to heatmap by taking mean over channels and taking absolute magnitude
        heat = integrated_grad[0].abs().mean(dim=0).cpu().numpy()
        mn, mx = heat.min(), heat.max()
        return ((heat - mn) / (mx - mn + 1e-8)).astype(np.float32)

    def _rise(self, tensor: torch.Tensor, target_class: Optional[int], N: int = 200, s: int = 8, p1: float = 0.5) -> np.ndarray:
        """RISE (Randomized Input Sampling for Explanation)."""
        with torch.no_grad():
            logits = self.model(tensor)
        if target_class is None:
            target_class = int(logits.argmax(-1))
            
        cell_size = int(np.ceil(224 / s))
        up_size = int((s + 1) * cell_size)
        
        np.random.seed(42)
        grid = np.random.rand(N, s, s) < p1
        grid = grid.astype(np.float32)
        
        masks = np.empty((N, 224, 224), dtype=np.float32)
        zoom_factor = up_size / s
        for i in range(N):
            m = scipy.ndimage.zoom(grid[i], (zoom_factor, zoom_factor), order=1)
            shift_x = np.random.randint(0, cell_size)
            shift_y = np.random.randint(0, cell_size)
            masks[i] = m[shift_y:shift_y+224, shift_x:shift_x+224]
            
        masks_t = torch.from_numpy(masks).unsqueeze(1).to(self.device) # (N, 1, 224, 224)
        attribution = np.zeros((224, 224), dtype=np.float32)
        
        batch_size = 16
        with torch.no_grad():
            for i in range(0, N, batch_size):
                b_masks = masks_t[i:i+batch_size]
                b_inputs = tensor.repeat(b_masks.shape[0], 1, 1, 1) * b_masks
                
                b_logits = self.model(b_inputs)
                b_probs = torch.softmax(b_logits, dim=-1)[:, target_class].cpu().numpy()
                
                for j in range(b_masks.shape[0]):
                    attribution += b_probs[j] * masks[i+j]
                    
        mn, mx = attribution.min(), attribution.max()
        return ((attribution - mn) / (mx - mn + 1e-8)).astype(np.float32)

    def _attention_rollout(self, tensor: torch.Tensor) -> np.ndarray:
        """Multiply attention matrices across all ViT layers."""
        with torch.no_grad():
            _ = self.model(tensor)
        attn_maps = self.model.get_attention_maps()  # list of (1, H, N, N) from vit.blocks
        if not attn_maps:
            return np.zeros((224, 224), dtype=np.float32)

        rollout = torch.eye(attn_maps[0].shape[-1], device=self.device)
        for attn in attn_maps:
            a = attn[0].mean(0)  # (N, N) – average over heads
            a = a + torch.eye(a.shape[0], device=self.device)
            a = a / a.sum(-1, keepdim=True)
            rollout = a @ rollout

        # CLS → patches
        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)

    def _occlusion_sensitivity(self, image: Image.Image,
                                target_class: Optional[int], grid: int = 8) -> np.ndarray:
        """Slide an occluded patch across the image and record score drop."""
        img_arr = np.array(image.resize((224, 224)).convert("RGB"), dtype=np.uint8)
        step = 224 // grid
        score_map = np.zeros((grid, grid), dtype=np.float32)

        tensor_orig = _to_tensor(image, self.device)
        with torch.no_grad():
            logits = self.model(tensor_orig)
            if target_class is None:
                target_class = int(logits.argmax(-1))
            base_score = float(F.softmax(logits, dim=-1)[0, target_class])

        mean_rgb = [int(m * 255) for m in IMAGENET_MEAN]
        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 = _to_tensor(occ_img, self.device)
                    s = float(F.softmax(self.model(t), dim=-1)[0, target_class])
                score_map[gy, gx] = base_score - s

        mn, mx = score_map.min(), score_map.max()
        return ((score_map - mn) / (mx - mn + 1e-8)).astype(np.float32)

    # ── 5. generate_class_comparison ──────────────────────────────────────────

    def generate_class_comparison(self) -> str:
        """
        4 rows (one per class) × 6 columns:
        Original, ViT3, ViT7, ViT11, Mamba0, Mamba3
        Color-coded row titles.
        Returns base64 PNG.
        """
        col_keys    = ["original", "vit_3", "vit_7", "vit_11", "mamba_0", "mamba_1"]
        col_labels  = ["Original", "ViT Blk 3", "ViT Blk 7", "ViT Blk 11",
                       "MedMamba 0", "MedMamba 1"]
        row_colors  = {
            "CNV":    "#FF4B4B",
            "DME":    "#4B7BFF",
            "DRUSEN": "#FFD700",
            "NORMAL": "#4BFF6F",
        }

        fig, axes = plt.subplots(4, 6, figsize=(18, 12), dpi=100)
        fig.patch.set_facecolor(BG_COLOR)
        fig.suptitle("RetViM Class Comparison – Feature Maps per Stage",
                     color="white", fontsize=13, fontweight="bold", y=0.98)

        for ri, cls in enumerate(CLASS_NAMES):
            syn = draw_synthetic_oct(cls)
            tensor = _to_tensor(syn, self.device)
            with torch.no_grad():
                feats = self.model.get_intermediate_features(tensor)

            oct_rgb = np.array(syn.resize((224, 224)).convert("RGB"))
            tint    = row_colors[cls]

            for ci, key in enumerate(col_keys):
                ax = axes[ri, ci]
                ax.set_facecolor("#161B22")
                if key == "original":
                    ax.imshow(oct_rgb)
                elif key in feats:
                    cmap = "plasma" if key.startswith("vit") else "inferno"
                    ax.imshow(_feature_to_map(feats[key]), cmap=cmap, interpolation="bilinear")
                ax.axis("off")

                if ci == 0:
                    ax.text(-0.08, 0.5, cls, transform=ax.transAxes,
                            color=tint, fontsize=10, fontweight="bold",
                            va="center", ha="right", rotation=90)
                if ri == 0:
                    ax.set_title(col_labels[ci], color="white", fontsize=8, pad=3)

        plt.tight_layout(pad=1.2)
        return _fig_to_b64(fig)

    # ── 6. generate_pca_plot ───────────────────────────────────────────────────

    def generate_pca_plot(self, samples_per_class: int = 5) -> str:
        """
        Extract features at 8 stages, apply PCA, show scatter plots.
        Returns base64 PNG.
        """
        stage_keys   = ["patch_embed", "vit_0", "vit_3", "vit_7", "vit_11",
                        "mamba_0", "mamba_1"]
        stage_labels = ["Patch Embed", "ViT Block 0", "ViT Block 3", "ViT Block 7",
                        "ViT Block 11", "MedMamba Block 0", "MedMamba Block 1"]
        colors = {"CNV": "#FF4B4B", "DME": "#4B7BFF", "DRUSEN": "#FFD700", "NORMAL": "#4BFF6F"}

        # Collect features
        all_feats  = {k: [] for k in stage_keys}
        all_labels = []

        np.random.seed(0)
        for cls in CLASS_NAMES:
            for _ in range(samples_per_class):
                syn = draw_synthetic_oct(cls)
                t = _to_tensor(syn, self.device)
                with torch.no_grad():
                    f = self.model.get_intermediate_features(t)
                for k in stage_keys:
                    if k in f:
                        vec = f[k][0].mean(0).cpu().numpy()  # (D,)
                        all_feats[k].append(vec)
                all_labels.append(cls)

        labels_arr = np.array(all_labels)

        fig, axes = plt.subplots(2, 4, figsize=(20, 10), dpi=100)
        fig.patch.set_facecolor(BG_COLOR)
        fig.suptitle("RetViM PCA Feature Embeddings", color="white",
                     fontsize=13, fontweight="bold")

        for idx, (key, label) in enumerate(zip(stage_keys, stage_labels)):
            ax = axes[idx // 4, idx % 4]
            ax.set_facecolor("#161B22")
            for sp in ax.spines.values():
                sp.set_color("#30363D")
            ax.tick_params(colors="white", labelsize=6)

            X = np.array(all_feats[key])
            if X.shape[0] >= 2:
                n_comp = min(2, X.shape[1], X.shape[0])
                pca = PCA(n_components=n_comp)
                Z   = pca.fit_transform(X)
                if Z.shape[1] == 1:
                    Z = np.hstack([Z, np.zeros_like(Z)])
                for cls in CLASS_NAMES:
                    mask = labels_arr == cls
                    ax.scatter(Z[mask, 0], Z[mask, 1],
                               c=colors[cls], label=cls, s=30, alpha=0.8)

            ax.set_title(label, color="white", fontsize=8, pad=4)
            ax.set_xlabel("PC1", color="white", fontsize=6)
            ax.set_ylabel("PC2", color="white", fontsize=6)
            if idx == 0:
                ax.legend(fontsize=6, facecolor="#161B22", labelcolor="white",
                          loc="upper right")

        plt.tight_layout(pad=1.5)
        return _fig_to_b64(fig)

    # ── 7. extract_deep_analysis ────────────────────────────────────────────────

    def extract_deep_analysis(self, image: Image.Image) -> dict:
        """
        Extract real computed data for the Mamba and Layers tabs.
        Returns attention heads, layer stats, Mamba internals — all from
        the actual model processing the actual uploaded image.
        """
        tensor = _to_tensor(image, self.device)

        # ── Enable Mamba internals storage ──
        self.model.enable_mamba_store(True)

        with torch.no_grad():
            # Single forward pass: populates attention maps + mamba internals
            _ = self.model(tensor)

        # ── 1. Attention heads: real per-head CLS→patch maps ──
        attn_maps = self.model.get_attention_maps()  # list of (1, 12, 197, 197)
        heads_data = []      # [12 layers][12 heads][14][14]
        head_entropy = []    # [12 layers][12 heads]
        layer_entropy = []   # [12 layers]

        for layer_attn in attn_maps:
            a = layer_attn[0]  # (12, 197, 197)
            layer_heads = []
            layer_head_ent = []
            layer_ent_sum = 0.0
            for h in range(a.shape[0]):
                # CLS → patch attention
                cls_to_patch = a[h, 0, 1:].cpu().numpy()  # (196,)
                side = int(round(cls_to_patch.shape[0] ** 0.5))
                hmap = cls_to_patch[:side*side].reshape(side, side)
                mn, mx = hmap.min(), hmap.max()
                hmap = ((hmap - mn) / (mx - mn + 1e-8)).astype(np.float32)
                layer_heads.append(hmap.tolist())

                # Per-head entropy from full attention distribution
                p = a[h].cpu().numpy().flatten()
                p = p[p > 1e-10]
                ent = float(-np.sum(p * np.log2(p + 1e-12)))
                layer_head_ent.append(round(ent, 4))
                layer_ent_sum += ent

            heads_data.append(layer_heads)
            head_entropy.append(layer_head_ent)
            layer_entropy.append(round(layer_ent_sum / a.shape[0], 4))

        # ── 2. Layer stats: magnitudes + CLS tokens ──
        with torch.no_grad():
            layer_info = self.model.get_all_layer_features(tensor)

        magnitudes = [round(m, 6) for m in layer_info["magnitudes"]]

        # ── 3. CLS class similarity ──
        cls_tokens = layer_info["cls_tokens"]  # list of 14 tensors (768,)
        cls_similarity = self._compute_cls_similarity(cls_tokens)

        # ── 4. Activation distributions (frozen vs trainable) ──
        frozen_acts = []
        trainable_acts = []
        for i, blk in enumerate(self.model.vit.blocks):
            if blk.attn.last_attn is not None:
                vals = blk.attn.last_attn[0].cpu().numpy().flatten()
                if i < 6:
                    frozen_acts.extend(vals[::10].tolist())
                else:
                    trainable_acts.extend(vals[::10].tolist())

        frozen_hist, frozen_edges = np.histogram(
            frozen_acts, bins=50, density=True
        ) if frozen_acts else (np.zeros(50), np.linspace(0, 1, 51))
        trainable_hist, trainable_edges = np.histogram(
            trainable_acts, bins=50, density=True
        ) if trainable_acts else (np.zeros(50), np.linspace(0, 1, 51))

        # ── 5. Mamba internals ──
        mamba_internals = self.model.get_mamba_internals()

        # Disable storage
        self.model.enable_mamba_store(False)

        return {
            "attention_heads":   heads_data,       # [12][12][14][14]
            "head_entropy":      head_entropy,     # [12][12]
            "layer_entropy":     layer_entropy,    # [12]
            "layer_magnitude":   magnitudes,       # [14]
            "cls_similarity":    cls_similarity,   # {class: [14 floats]}
            "frozen_hist":       frozen_hist.tolist(),
            "trainable_hist":    trainable_hist.tolist(),
            "hist_bins":         frozen_edges.tolist(),
            "mamba_internals":   mamba_internals,   # [2 blocks]
        }

    def _compute_cls_similarity(self, cls_tokens: list) -> dict:
        """Compute cosine similarity of CLS token at each layer to class prototypes."""
        # Build class prototypes (lazy, cached)
        if not hasattr(self, '_class_prototypes'):
            self._class_prototypes = {}
            for cls_name in CLASS_NAMES:
                syn = draw_synthetic_oct(cls_name)
                t = _to_tensor(syn, self.device)
                with torch.no_grad():
                    logits = self.model(t)
                # Get final CLS token after full forward
                # Re-run to get the CLS from the last mamba block
                B = t.shape[0]
                tokens = self.model.vit.patch_embed(t)
                cls = self.model.vit.cls_token.expand(B, -1, -1)
                tokens = torch.cat([cls, tokens], dim=1)
                tokens = self.model.vit.pos_drop(tokens + self.model.vit.pos_embed)
                for blk in self.model.vit.blocks:
                    tokens = blk(tokens)
                tokens = self.model.vit.norm(tokens)
                for blk in self.model.medmamba_blocks:
                    tokens = blk(tokens)
                self._class_prototypes[cls_name] = tokens[0, 0].detach().cpu()

        result = {}
        for cls_name, proto in self._class_prototypes.items():
            sims = []
            for ct in cls_tokens:
                cos = float(F.cosine_similarity(ct.unsqueeze(0), proto.unsqueeze(0)))
                sims.append(round(cos, 6))
            result[cls_name] = sims
        return result

    # ── 8. full_analysis ───────────────────────────────────────────────────────

    def full_analysis(self, image: Image.Image, fast: bool = False) -> dict:
        """
        Run all visualisations. Cache result to /tmp for 10 minutes.
        fast=True: skips RISE and IG (for live uploaded images, target <60s).
        fast=False: full suite including RISE and IG (for pre-generated samples).
        Returns complete dict with all base64 images.
        """
        # Compute cache key from image pixels (include fast flag in key)
        img_bytes = image.tobytes()
        key = hashlib.md5(img_bytes).hexdigest() + ("_fast" if fast else "_full")

        cached = _load_cache(key)
        if cached is not None and "deep_analysis" in cached:
            return cached

        pred   = self.predict(image)
        cls_nm = pred["prediction"]

        journey  = self.generate_neural_journey(image, cls_nm)
        gradcam  = self.generate_gradcam_fast(image) if fast else self.generate_gradcam(image)
        cls_cmp  = self.generate_class_comparison()
        pca      = self.generate_pca_plot()
        deep     = self.extract_deep_analysis(image)

        result = {
            "prediction":         pred,
            "neural_journey_b64": journey,
            "gcam_b64":           gradcam["gcam_b64"],
            "gcam2_b64":          gradcam.get("gcam2_b64"),
            "rollout_b64":        gradcam.get("rollout_b64"),
            "occ_b64":            gradcam.get("occ_b64"),
            "ig_b64":             gradcam.get("ig_b64"),
            "rise_b64":           gradcam.get("rise_b64"),
            "overlay_b64":        gradcam.get("overlay_b64"),
            "class_comparison_b64": cls_cmp,
            "pca_b64":            pca,
            "deep_analysis":      deep,
        }

        _save_cache(key, result)
        return result


# ── legacy helpers (kept for backward compat) ──────────────────────────────────

def preprocess_image(image: Image.Image) -> torch.Tensor:
    if image.mode != "RGB":
        image = image.convert("RGB")
    return _transform(image).unsqueeze(0)


def pil_to_b64(image: Image.Image, fmt: str = "PNG") -> str:
    return _pil_to_b64(image)


def predict(model, image, class_names, device, top_k=4, generate_cam=True) -> dict:
    """Legacy functional predict (used by old main.py routes)."""
    import torch.nn.functional as F
    tensor = preprocess_image(image).to(device)
    with torch.no_grad():
        logits = model(tensor)
    probs = F.softmax(logits, dim=-1)[0].cpu().tolist()
    pred_idx = int(np.argmax(probs))
    pred_class = class_names[pred_idx]
    confidence = probs[pred_idx]
    top_pairs = sorted(zip(class_names, probs), key=lambda x: -x[1])[:top_k]
    return {
        "predicted_class": pred_class,
        "confidence": round(confidence, 6),
        "probabilities": {c: round(p, 6) for c, p in zip(class_names, probs)},
        "top_k": [{"class": c, "probability": round(p, 6)} for c, p in top_pairs],
        "gradcam_b64": None,
        "chart_b64": "",
    }