gradcam_v3.py

#!/usr/bin/env python3
"""
RetinaSense v3.0 — Grad-CAM Explainability Pipeline
====================================================
Implements:
  1. ViTGradCAM   : Gradient-weighted Class Activation Maps for ViT backbone
  2. OODDetector  : Mahalanobis-distance out-of-distribution detection
  3. predict_with_gradcam : Full inference pipeline (preprocess → OOD → CAM → calibrate)
  4. Batch evaluation on 20 test images (4 per class)
  5. Disease-specific heatmap validation against known anatomical regions
  6. Clinical output report (GRADCAM_REPORT.md)

Usage:
  python gradcam_v3.py
"""

import os
import sys
import json
import warnings
import numpy as np
import cv2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from PIL import Image
from datetime import datetime
import time

warnings.filterwarnings('ignore')

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

import timm

# ================================================================
# CONFIGURATION
# ================================================================
BASE_DIR   = '/teamspace/studios/this_studio'
OUTPUT_DIR = os.path.join(BASE_DIR, 'outputs_v3')
GRADCAM_DIR = os.path.join(OUTPUT_DIR, 'gradcam')
os.makedirs(GRADCAM_DIR, exist_ok=True)

MODEL_PATH      = os.path.join(OUTPUT_DIR, 'best_model.pth')
THRESHOLDS_PATH = os.path.join(OUTPUT_DIR, 'thresholds.json')
TEMPERATURE_PATH = os.path.join(OUTPUT_DIR, 'temperature.json')
TEST_CSV        = os.path.join(BASE_DIR, 'data', 'test_split.csv')
NORM_STATS_PATH = os.path.join(BASE_DIR, 'data', 'fundus_norm_stats.json')

CLASS_NAMES = ['Normal', 'Diabetes/DR', 'Glaucoma', 'Cataract', 'AMD']
NUM_CLASSES = 5
IMG_SIZE    = 224
DROPOUT     = 0.3

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Anatomical regions expected for each disease class
EXPECTED_REGIONS = {
    0: 'low uniform activation (Normal)',
    1: 'scattered periphery and macula (DR)',
    2: 'optic disc (Glaucoma)',
    3: 'diffuse lens opacity (Cataract)',
    4: 'macula/centre-temporal (AMD)',
}

print('=' * 65)
print('   RetinaSense v3.0 — Grad-CAM Explainability Pipeline')
print('=' * 65)
print(f'  Device   : {DEVICE}')
if torch.cuda.is_available():
    print(f'  GPU      : {torch.cuda.get_device_name(0)}')
print(f'  Output   : {GRADCAM_DIR}')
print('=' * 65)


# ================================================================
# LOAD NORMALISATION STATS
# ================================================================
if os.path.exists(NORM_STATS_PATH):
    with open(NORM_STATS_PATH) as f:
        norm_stats = json.load(f)
    NORM_MEAN = norm_stats['mean_rgb']
    NORM_STD  = norm_stats['std_rgb']
    print(f'  Fundus norm stats: mean={[round(v,4) for v in NORM_MEAN]}, std={[round(v,4) for v in NORM_STD]}')
else:
    NORM_MEAN = [0.485, 0.456, 0.406]
    NORM_STD  = [0.229, 0.224, 0.225]
    print('  Using ImageNet normalisation fallback')


# ================================================================
# MODEL ARCHITECTURE  (mirrors retinasense_v3.py exactly)
# ================================================================
class MultiTaskViT(nn.Module):
    """ViT-Base-Patch16-224 with disease + severity heads."""

    def __init__(self, n_disease=NUM_CLASSES, n_severity=5, drop=DROPOUT):
        super().__init__()
        self.backbone = timm.create_model(
            'vit_base_patch16_224', pretrained=False, num_classes=0
        )
        feat = 768  # CLS token dimension

        self.drop = nn.Dropout(drop)

        self.disease_head = nn.Sequential(
            nn.Linear(feat, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(512, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, n_disease),
        )
        self.severity_head = nn.Sequential(
            nn.Linear(feat, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(256, n_severity),
        )

    def forward(self, x):
        f = self.backbone(x)   # (B, 768) — CLS token features
        f = self.drop(f)
        return self.disease_head(f), self.severity_head(f)

    def get_features(self, x):
        """Return raw CLS token features (before heads and dropout)."""
        return self.backbone(x)  # (B, 768)

    def forward_with_tokens(self, x):
        """Return (disease_logits, full_token_sequence (B,197,768))."""
        tokens = self.backbone.forward_features(x)  # (B, 197, 768)
        cls_feat = tokens[:, 0, :]
        cls_feat_d = self.drop(cls_feat)
        d_out = self.disease_head(cls_feat_d)
        return d_out, tokens


# ================================================================
# LOAD MODEL
# ================================================================
print('\nLoading model...')
model = MultiTaskViT().to(DEVICE)
ckpt = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
model.load_state_dict(ckpt['model_state_dict'])
model.eval()
print(f'  Loaded: {MODEL_PATH}')
print(f'  Checkpoint epoch: {ckpt.get("epoch", "?") + 1}  val_acc={ckpt.get("val_acc", 0):.2f}%')

# Load thresholds and temperature
with open(THRESHOLDS_PATH) as f:
    thr_data = json.load(f)
THRESHOLDS = thr_data['thresholds']

with open(TEMPERATURE_PATH) as f:
    temp_data = json.load(f)
TEMPERATURE = temp_data['temperature']

print(f'  Temperature T = {TEMPERATURE:.4f}')
print(f'  Thresholds    = {[round(t,3) for t in THRESHOLDS]}')


# ================================================================
# IMAGE PREPROCESSING
# ================================================================
def ben_graham(path, sz=IMG_SIZE, sigma=10):
    """Ben Graham high-frequency fundus enhancement (APTOS-style)."""
    img = cv2.imread(path)
    if img is None:
        img = np.array(Image.open(path).convert('RGB'))
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (sz, sz))
    img = cv2.addWeighted(img, 4, cv2.GaussianBlur(img, (0, 0), sigma), -4, 128)
    mask = np.zeros(img.shape[:2], dtype=np.uint8)
    cv2.circle(mask, (sz // 2, sz // 2), int(sz * 0.48), 255, -1)
    return cv2.bitwise_and(img, img, mask=mask)


def clahe_preprocess(path, sz=IMG_SIZE):
    """CLAHE-based contrast enhancement (ODIR-style)."""
    img = cv2.imread(path)
    if img is None:
        img = np.array(Image.open(path).convert('RGB'))
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    img = cv2.resize(img, (sz, sz))
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    lab[:, :, 0] = clahe.apply(lab[:, :, 0])
    img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


def load_and_preprocess(image_path, dataset='auto'):
    """
    Load image and apply domain-conditional preprocessing.
    Returns:
        img_np   : numpy (224,224,3) uint8 preprocessed
        img_orig : numpy (224,224,3) uint8 original (for overlay)
    """
    # Normalise path: handle relative paths from CSV (e.g. "aptos/..." or "./aptos/...")
    # If the path is already absolute and exists, use it directly.
    # Otherwise resolve relative to BASE_DIR, stripping any leading ./ or .// first.
    if not os.path.isabs(image_path):
        # Strip any leading './' or '../' patterns to get a clean relative path
        clean = image_path
        while clean.startswith('./') or clean.startswith('.//'):
            clean = clean[2:] if clean.startswith('./') else clean[3:]
        image_path = os.path.join(BASE_DIR, clean)thinl
    # Auto-detect domain
    if dataset == 'auto':
        if 'aptos' in image_path.lower() or 'gaussian' in image_path.lower():
            dataset = 'APTOS'
        else:
            dataset = 'ODIR'

    # Load original (unprocessed, for overlay)
    raw = cv2.imread(image_path)
    if raw is None:
        raw = np.array(Image.open(image_path).convert('RGB'))
    else:
        raw = cv2.cvtColor(raw, cv2.COLOR_BGR2RGB)
    img_orig = cv2.resize(raw, (IMG_SIZE, IMG_SIZE))

    # Apply preprocessing
    if dataset == 'APTOS':
        img_np = ben_graham(image_path)
    else:
        img_np = clahe_preprocess(image_path)

    return img_np, img_orig


def preprocess_to_tensor(img_np):
    """Convert preprocessed numpy image to normalised tensor (1, 3, 224, 224)."""
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.ToTensor(),
        transforms.Normalize(NORM_MEAN, NORM_STD),
    ])
    return transform(img_np).unsqueeze(0)


# ================================================================
# ViT GRAD-CAM
# ================================================================
class ViTAttentionRollout:
    """
    Attention Rollout for Vision Transformer (Abnar & Zuidema, 2020).

    WHY this works better than Grad-CAM for ViT:
    - ViT uses CLS token pooling: gradients flow ONLY through CLS token (index 0)
    - All 196 patch token gradients at block output = zero → Grad-CAM fails
    - Attention Rollout instead traces how information flows from image patches
      to the CLS token across ALL 12 transformer layers
    - Accounts for residual connections by adding identity to each attention map
    - Produces spatially meaningful maps that highlight actual disease regions

    Algorithm:
      1. Collect attention maps A_l from all 12 blocks: shape (B, H, N, N)
      2. Average over H heads: A_l → (B, N, N)
      3. Add identity: A_l = A_l + I  (accounts for residual connection)
      4. Row-normalize: A_l = A_l / row_sum
      5. Matrix-multiply all layers: Rollout = A_0 @ A_1 @ ... @ A_11
      6. Take CLS row, patch tokens only: Rollout[0, 1:] → (196,)
      7. Reshape 14×14 → bilinear upsample → 224×224
    """

    def __init__(self, model, discard_ratio=0.97):
        self.model = model
        self.discard_ratio = discard_ratio  # zero out weakest attention weights
        self._attention_maps = []
        self._hooks = []

        # Disable fused attention for explicit weight access
        for block in model.backbone.blocks:
            block.attn.fused_attn = False

        # Register forward hooks on ALL transformer blocks
        for block in model.backbone.blocks:
            h = block.attn.register_forward_hook(self._attn_hook)
            self._hooks.append(h)

    def _attn_hook(self, module, input, output):
        """Capture softmax attention weights from each block."""
        x = input[0]
        B, N, C = x.shape
        with torch.no_grad():
            qkv = module.qkv(x).reshape(B, N, 3, module.num_heads, module.head_dim).permute(2, 0, 3, 1, 4)
            q, k, _ = qkv.unbind(0)
            q, k = module.q_norm(q), module.k_norm(k)
            attn = (q * module.scale @ k.transpose(-2, -1)).softmax(dim=-1)
            self._attention_maps.append(attn.detach().cpu())  # (B, H, N, N)

    def generate(self, image_tensor, class_idx=None):
        """
        Generate attention rollout heatmap.

        Returns:
            heatmap         : np.ndarray (224, 224) float32 [0, 1]
                              High values = regions most important for prediction
            predicted_label : int
            confidence      : float (raw softmax)
        """
        self.model.eval()
        self._attention_maps = []

        with torch.no_grad():
            image_tensor = image_tensor.to(DEVICE)
            d_out, _ = self.model(image_tensor)
            probs = torch.softmax(d_out, dim=1)
            predicted_label = int(probs.argmax(dim=1).item())
            confidence = float(probs[0, predicted_label].item())

        if class_idx is None:
            class_idx = predicted_label

        # --- Attention Rollout computation ---
        # Stack all layer attentions: list of (1, H, N, N) → (L, H, N, N)
        attn_stack = torch.stack(self._attention_maps, dim=0)  # (L, 1, H, N, N)
        attn_stack = attn_stack[:, 0]  # (L, H, N, N), batch=1

        # Average over heads
        attn_mean = attn_stack.mean(dim=1)  # (L, N, N)

        # Optional: discard weakest connections (sharpens the map)
        if self.discard_ratio > 0:
            flat = attn_mean.reshape(attn_mean.shape[0], -1)
            thresh = torch.quantile(flat, self.discard_ratio, dim=1, keepdim=True)
            thresh = thresh.unsqueeze(-1)  # broadcast over N,N
            attn_mean = torch.where(attn_mean >= thresh, attn_mean, torch.zeros_like(attn_mean))

        # Add identity matrix for residual connection, then row-normalize
        I = torch.eye(attn_mean.shape[-1]).unsqueeze(0)  # (1, N, N)
        attn_aug = attn_mean + I
        attn_aug = attn_aug / attn_aug.sum(dim=-1, keepdim=True).clamp(min=1e-8)

        # Matrix-multiply across all layers
        rollout = attn_aug[0]
        for l in range(1, len(attn_aug)):
            rollout = rollout @ attn_aug[l]

        # CLS token's attention to all patch tokens (skip CLS at index 0)
        cls_attention = rollout[0, 1:]  # (196,)

        # Reshape and upsample
        spatial = cls_attention.numpy().reshape(14, 14).astype(np.float32)
        spatial = cv2.resize(spatial, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_LINEAR)

        # Normalize to [0, 1]
        s_min, s_max = spatial.min(), spatial.max()
        if s_max - s_min > 1e-8:
            spatial = (spatial - s_min) / (s_max - s_min)
        else:
            spatial = np.zeros_like(spatial)

        # Power-curve stretch: boosts mid-range attention values for visual clarity
        # gamma < 1 brightens the map; 0.4 gives strong contrast enhancement
        spatial = np.power(spatial, 0.4)

        return spatial.astype(np.float32), predicted_label, confidence

    def overlay(self, original_image_np, heatmap, alpha=0.5):
        """
        Blend attention rollout heatmap onto original fundus image.
        Uses INFERNO colormap (dark=low, bright=high) — better for medical images.

        Args:
            original_image_np : (224, 224, 3) uint8 RGB
            heatmap           : (224, 224) float32 [0, 1]
            alpha             : heatmap opacity (0.5 gives good visibility)

        Returns:
            overlay : (224, 224, 3) uint8 RGB
        """
        # Apply JET colormap
        heatmap_uint8 = (heatmap * 255).astype(np.uint8)
        colormap = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
        colormap_rgb = cv2.cvtColor(colormap, cv2.COLOR_BGR2RGB)

        # Apply circular mask to ignore black borders (fundus images are circular)
        h, w = heatmap.shape
        cy, cx = h // 2, w // 2
        radius = min(h, w) // 2 - 5
        mask = np.zeros((h, w), dtype=np.float32)
        cv2.circle(mask, (cx, cy), radius, 1.0, -1)
        mask = cv2.GaussianBlur(mask, (21, 21), 0)

        # Blend only inside the retinal circle
        orig = original_image_np.astype(np.float32)
        cmap = colormap_rgb.astype(np.float32)
        blended = orig.copy()
        for c in range(3):
            blended[:, :, c] = (
                orig[:, :, c] * (1 - alpha * mask)
                + cmap[:, :, c] * (alpha * mask)
            )
        return np.clip(blended, 0, 255).astype(np.uint8)

    def remove_hooks(self):
        """Clean up all registered hooks."""
        for h in self._hooks:
            h.remove()
        self._hooks = []

# Keep old name as alias for backward compatibility
ViTGradCAM = ViTAttentionRollout



# ================================================================
# OOD DETECTION (Mahalanobis Distance)
# ================================================================
class OODDetector:
    """
    Out-of-Distribution detector using class-conditional Mahalanobis distance.

    Fit on training-set CLS token features; at inference, computes the
    minimum Mahalanobis distance from the test feature to the nearest
    class centroid. High distance = likely OOD.
    """

    def __init__(self, threshold_percentile=97.5):
        self.class_means   = None   # (num_classes, feat_dim)
        self.cov_inv       = None   # (feat_dim, feat_dim)
        self.ood_threshold = None
        self.threshold_percentile = threshold_percentile
        self.is_fitted     = False

    def fit(self, model, dataloader, device, max_batches=60):
        """
        Extract CLS token features for all samples, compute class-conditional
        means and shared inverse covariance matrix.
        """
        print('  OODDetector.fit: extracting features...')
        all_features = []
        all_labels   = []

        model.eval()
        with torch.no_grad():
            for i, batch in enumerate(dataloader):
                if i >= max_batches:
                    break
                imgs, d_lbl, _ = batch
                imgs = imgs.to(device)
                feats = model.get_features(imgs)  # (B, 768)
                all_features.append(feats.cpu().numpy())
                all_labels.append(d_lbl.numpy())

        features = np.concatenate(all_features, axis=0)  # (N, 768)
        labels   = np.concatenate(all_labels,   axis=0)  # (N,)

        num_classes = NUM_CLASSES
        feat_dim    = features.shape[1]

        # Class-conditional means
        self.class_means = np.zeros((num_classes, feat_dim), dtype=np.float64)
        for c in range(num_classes):
            mask = labels == c
            if mask.sum() > 0:
                self.class_means[c] = features[mask].mean(axis=0)

        # Shared (pooled) covariance matrix
        cov = np.zeros((feat_dim, feat_dim), dtype=np.float64)
        total = 0
        for c in range(num_classes):
            mask = labels == c
            if mask.sum() < 2:
                continue
            diff = features[mask] - self.class_means[c]
            cov += diff.T @ diff
            total += mask.sum()

        cov /= max(total - num_classes, 1)

        # Regularise for numerical stability (add small diagonal)
        cov += np.eye(feat_dim) * 1e-4

        # Pseudo-inverse via SVD (numerically stable for high-dim)
        try:
            self.cov_inv = np.linalg.pinv(cov)
        except np.linalg.LinAlgError:
            self.cov_inv = np.eye(feat_dim)

        # Compute train-set Mahalanobis distances to set threshold
        train_dists = []
        for feat in features:
            d = self._mahal_min_dist(feat)
            train_dists.append(d)
        self.ood_threshold = float(np.percentile(train_dists, self.threshold_percentile))

        self.is_fitted = True
        print(f'  OOD threshold ({self.threshold_percentile}th pct): {self.ood_threshold:.4f}')
        print(f'  Features extracted: {len(features)} samples')

    def _mahal_min_dist(self, feat):
        """Minimum Mahalanobis distance to any class centroid."""
        min_dist = float('inf')
        for c in range(NUM_CLASSES):
            diff = feat - self.class_means[c]
            dist = float(diff @ self.cov_inv @ diff)
            dist = max(dist, 0.0)  # guard against floating-point negatives
            if dist < min_dist:
                min_dist = dist
        return np.sqrt(min_dist)

    def score(self, features):
        """
        Compute OOD score for a batch of features.

        Args:
            features : np.ndarray (N, 768) or (768,)

        Returns:
            distances : np.ndarray (N,) Mahalanobis distances
            ood_flags : np.ndarray (N,) bool, True = likely OOD
        """
        if not self.is_fitted:
            raise RuntimeError('OODDetector.fit() must be called before score()')

        if features.ndim == 1:
            features = features[np.newaxis, :]

        distances = np.array([self._mahal_min_dist(f) for f in features])
        ood_flags = distances > self.ood_threshold
        return distances, ood_flags

    def save(self, path):
        np.savez(path,
                 class_means=self.class_means,
                 cov_inv=self.cov_inv,
                 ood_threshold=np.array([self.ood_threshold]),
                 threshold_percentile=np.array([self.threshold_percentile]))
        print(f'  OOD detector saved -> {path}.npz')

    def load(self, path):
        if not path.endswith('.npz'):
            path = path + '.npz'
        data = np.load(path)
        self.class_means           = data['class_means']
        self.cov_inv               = data['cov_inv']
        self.ood_threshold         = float(data['ood_threshold'][0])
        self.threshold_percentile  = float(data['threshold_percentile'][0])
        self.is_fitted             = True
        print(f'  OOD detector loaded <- {path}')


# ================================================================
# ATTENTION REGION ANALYSER
# ================================================================
def analyse_attention_region(heatmap, disease_class):
    """
    Check if the Grad-CAM heatmap activation pattern is consistent
    with the expected anatomical region for the given disease.

    Returns:
        attention_region : str describing where activation is
        is_consistent   : bool
        region_scores   : dict with activation energy in each zone
    """
    h, w = heatmap.shape  # (224, 224)
    cx, cy = w // 2, h // 2

    # Define anatomical zones (approximate, relative to image centre)
    # Centre disc zone: circle r ~ 30px (optic disc)
    r_disc   = int(h * 0.13)
    # Macula zone: circle r ~ 55px centred slightly temporal
    r_macula = int(h * 0.25)
    cx_mac   = int(cx + w * 0.10)  # slightly nasal offset

    # Build zone masks
    Y, X = np.ogrid[:h, :w]

    # Optic disc (small circle, centre of image)
    disc_mask = ((X - cx)**2 + (Y - cy)**2) <= r_disc**2

    # Macula (larger circle, centre-temporal)
    macula_mask = ((X - cx_mac)**2 + (Y - cy)**2) <= r_macula**2

    # Periphery: outer 30% of image
    peri_mask = (X < int(w * 0.15)) | (X > int(w * 0.85)) | \
                (Y < int(h * 0.15)) | (Y > int(h * 0.85))

    # Compute mean activation in each zone
    disc_score   = float(heatmap[disc_mask].mean())   if disc_mask.sum() > 0 else 0.0
    macula_score = float(heatmap[macula_mask].mean()) if macula_mask.sum() > 0 else 0.0
    peri_score   = float(heatmap[peri_mask].mean())   if peri_mask.sum() > 0 else 0.0
    overall_mean = float(heatmap.mean())

    region_scores = {
        'optic_disc': round(disc_score, 4),
        'macula':     round(macula_score, 4),
        'periphery':  round(peri_score, 4),
        'overall':    round(overall_mean, 4),
    }

    # Determine dominant region label
    max_zone = max(region_scores, key=lambda k: region_scores[k] if k != 'overall' else -1)

    zone_labels = {
        'optic_disc': 'optic disc (centre)',
        'macula':     'macula (centre-temporal)',
        'periphery':  'scattered periphery',
    }
    dominant_label = zone_labels.get(max_zone, 'diffuse')

    # Assess uniformity (low std = diffuse / uniform)
    if heatmap.std() < 0.10:
        dominant_label = 'diffuse (low activation)'

    # Check consistency with expected region
    consistency_map = {
        0: lambda s: s['overall'] < 0.25,                     # Normal → low uniform
        1: lambda s: s['periphery'] > 0.20 or s['macula'] > 0.25,  # DR → periphery/macula
        2: lambda s: s['optic_disc'] > 0.30,                   # Glaucoma → disc
        3: lambda s: heatmap.std() < 0.15,                     # Cataract → diffuse
        4: lambda s: s['macula'] > 0.25,                       # AMD → macula
    }
    check_fn = consistency_map.get(disease_class, lambda s: True)
    is_consistent = check_fn(region_scores)

    return dominant_label, is_consistent, region_scores


# ================================================================
# FULL INFERENCE PIPELINE
# ================================================================
def predict_with_gradcam(image_path, model, gradcam, ood_detector,
                         thresholds, temperature, device,
                         true_label=None, dataset='auto'):
    """
    End-to-end inference with Grad-CAM and OOD detection.

    Steps:
      1. Load and preprocess image (Ben Graham for APTOS, CLAHE for ODIR)
      2. OOD check on ViT CLS token features
      3. Generate Grad-CAM heatmap
      4. Apply temperature scaling to logits
      5. Apply per-class thresholds
      6. Analyse attention region

    Returns:
        dict with predicted_class, confidence, gradcam_heatmap, etc.
    """
    # 1. Preprocess
    img_np, img_orig = load_and_preprocess(image_path, dataset=dataset)
    img_tensor = preprocess_to_tensor(img_np).to(device)

    # 2. OOD check using raw CLS features
    model.eval()
    with torch.no_grad():
        features = model.get_features(img_tensor).cpu().numpy()  # (1, 768)

    if ood_detector.is_fitted:
        distances, ood_flags = ood_detector.score(features)
        ood_distance = float(distances[0])
        ood_flag     = bool(ood_flags[0])
    else:
        ood_distance = 0.0
        ood_flag     = False

    # 3. Generate Grad-CAM (also runs forward + backward pass)
    heatmap, predicted_label, raw_confidence = gradcam.generate(img_tensor)

    # 4. Temperature-scaled calibrated probabilities
    # Run a clean no-grad forward pass to get stable logits for calibration
    model.eval()
    with torch.no_grad():
        raw_feats = model.backbone(img_tensor)   # (1, 768)
        raw_feats = model.drop(raw_feats)
        logits = model.disease_head(raw_feats).float().cpu()  # (1, 5)

    scaled_logits = logits / temperature
    calibrated_probs = torch.softmax(scaled_logits, dim=1)[0].numpy()  # (5,)

    # 5. Apply per-class thresholds
    above = [i for i, (p, t) in enumerate(zip(calibrated_probs, thresholds)) if p >= t]
    if above:
        final_label = int(above[np.argmax([calibrated_probs[i] for i in above])])
    else:
        final_label = int(np.argmax(calibrated_probs))

    final_confidence = float(calibrated_probs[final_label])
    predicted_class  = CLASS_NAMES[final_label]

    # 6. Heatmap overlay
    gradcam_overlay = gradcam.overlay(img_orig, heatmap, alpha=0.7)

    # 7. Attention region analysis
    attention_region, region_consistent, region_scores = analyse_attention_region(
        heatmap, final_label
    )

    # Append disease name for clarity
    disease_tag = CLASS_NAMES[final_label].replace('/', '-')
    attention_region_full = f'{attention_region} ({disease_tag})'

    # 8. Review flag: low confidence OR OOD
    review_flag = ood_flag or final_confidence < 0.50

    return {
        'image_path':       image_path,
        'predicted_class':  predicted_class,
        'predicted_label':  final_label,
        'confidence':       round(final_confidence, 4),
        'raw_confidence':   round(raw_confidence, 4),
        'all_probabilities': [round(float(p), 4) for p in calibrated_probs],
        'gradcam_heatmap':  heatmap,           # (224, 224) float32
        'gradcam_overlay':  gradcam_overlay,   # (224, 224, 3) uint8
        'img_orig':         img_orig,          # original for display
        'ood_flag':         ood_flag,
        'ood_distance':     round(ood_distance, 4),
        'review_flag':      review_flag,
        'attention_region': attention_region_full,
        'region_scores':    region_scores,
        'region_consistent': region_consistent,
        'true_label':       true_label,
    }


# ================================================================
# BATCH EVALUATION
# ================================================================
def run_batch_evaluation(model, gradcam, ood_detector,
                         thresholds, temperature, device,
                         n_per_class=4):
    """
    Run inference on n_per_class images per disease class (20 total).
    Saves individual overlay images + summary grid.
    """
    import pandas as pd
    print(f'\nRunning batch evaluation ({n_per_class} per class = {n_per_class * NUM_CLASSES} total)...')

    df = pd.read_csv(TEST_CSV)

    # Collect n_per_class unique samples per class
    samples = []
    for label in range(NUM_CLASSES):
        subset = df[df['disease_label'] == label].drop_duplicates(subset='image_path')
        chosen = subset.head(n_per_class)
        for _, row in chosen.iterrows():
            samples.append({
                'image_path':  row['image_path'],
                'true_label':  int(row['disease_label']),
                'dataset':     str(row.get('dataset', 'auto')),
            })

    results = []
    failed  = []

    for i, sample in enumerate(samples):
        img_path   = sample['image_path']
        true_label = sample['true_label']
        dataset    = sample['dataset']

        print(f'  [{i+1:2d}/{len(samples)}] {CLASS_NAMES[true_label]:15s} | {os.path.basename(img_path)}', end='  ')

        try:
            result = predict_with_gradcam(
                img_path, model, gradcam, ood_detector,
                thresholds, temperature, device,
                true_label=true_label,
                dataset=dataset,
            )
            correct = (result['predicted_label'] == true_label)
            flag_str = ' [OOD]' if result['ood_flag'] else ''
            flag_str += ' [REVIEW]' if result['review_flag'] else ''
            print(f'-> pred={result["predicted_class"]:15s}  conf={result["confidence"]:.3f}  {"OK" if correct else "WRONG"}{flag_str}')

            # Save overlay image
            save_name = f'gradcam_{i+1:02d}_true{true_label}_pred{result["predicted_label"]}_{os.path.splitext(os.path.basename(img_path))[0][:20]}.png'
            save_path = os.path.join(GRADCAM_DIR, save_name)

            fig, axes = plt.subplots(1, 3, figsize=(12, 4))
            axes[0].imshow(result['img_orig'])
            axes[0].set_title(f'Original\nTrue: {CLASS_NAMES[true_label]}', fontsize=9)
            axes[0].axis('off')

            axes[1].imshow(result['gradcam_heatmap'], cmap='jet', vmin=0, vmax=1)
            axes[1].set_title('Grad-CAM Heatmap', fontsize=9)
            axes[1].axis('off')

            axes[2].imshow(result['gradcam_overlay'])
            flag_line = ' [OOD]' if result['ood_flag'] else ''
            axes[2].set_title(
                f'Overlay\nPred: {result["predicted_class"]} ({result["confidence"]:.2f}){flag_line}',
                fontsize=9, color='red' if not correct else 'green'
            )
            axes[2].axis('off')

            plt.suptitle(
                f'Attention: {result["attention_region"]}',
                fontsize=8, color='gray'
            )
            plt.tight_layout()
            plt.savefig(save_path, dpi=120, bbox_inches='tight')
            plt.close()

            result['save_path'] = save_path
            results.append(result)

        except Exception as e:
            print(f'  ERROR: {e}')
            failed.append({'image_path': img_path, 'error': str(e)})

    return results, failed


# ================================================================
# SUMMARY GRID  (4 rows = classes 0-4, 4 cols = samples)
# ================================================================
def save_summary_grid(results):
    """Save a 5×4 summary grid (rows=classes, cols=samples)."""
    n_rows = NUM_CLASSES
    n_cols = 4

    # Group results by true label
    by_class = {i: [] for i in range(NUM_CLASSES)}
    for r in results:
        tl = r.get('true_label', r['predicted_label'])
        by_class[tl].append(r)

    fig, axes = plt.subplots(n_rows, n_cols, figsize=(16, 20))
    fig.patch.set_facecolor('#1a1a2e')

    for row_idx in range(n_rows):
        class_results = by_class[row_idx]
        for col_idx in range(n_cols):
            ax = axes[row_idx, col_idx]
            if col_idx < len(class_results):
                r = class_results[col_idx]
                ax.imshow(r['gradcam_overlay'])
                correct = (r['predicted_label'] == r.get('true_label', r['predicted_label']))
                border_color = '#2ecc71' if correct else '#e74c3c'
                for spine in ax.spines.values():
                    spine.set_edgecolor(border_color)
                    spine.set_linewidth(3)
                label_str = f'{r["predicted_class"]}\n{r["confidence"]:.2f}'
                if r['ood_flag']:
                    label_str += '\n[OOD]'
                ax.set_title(label_str, fontsize=7, color='white', pad=2)
            else:
                ax.set_facecolor('#1a1a2e')

            ax.axis('off')
            if col_idx == 0:
                ax.set_ylabel(CLASS_NAMES[row_idx], rotation=0, labelpad=50,
                              fontsize=10, color='white', fontweight='bold',
                              va='center')

    plt.suptitle(
        'RetinaSense v3.0 — Grad-CAM Summary Grid\n'
        'Rows = True Class | Green border = Correct | Red border = Wrong',
        fontsize=12, color='white', y=1.01
    )
    plt.tight_layout()
    grid_path = os.path.join(GRADCAM_DIR, 'gradcam_summary_grid.png')
    plt.savefig(grid_path, dpi=130, bbox_inches='tight',
                facecolor=fig.get_facecolor())
    plt.close()
    print(f'  Summary grid saved -> {grid_path}')
    return grid_path


# ================================================================
# DISEASE-SPECIFIC HEATMAP VALIDATION
# ================================================================
def validate_heatmaps(results):
    """
    Check per-disease whether Grad-CAM activates the expected anatomical region.
    Returns a validation summary dict, saves to heatmap_validation.json.
    """
    print('\nRunning disease-specific heatmap validation...')

    validation = {}
    for cls_idx, cls_name in enumerate(CLASS_NAMES):
        cls_results = [r for r in results if r.get('true_label') == cls_idx]
        if not cls_results:
            validation[cls_name] = {'n_samples': 0}
            continue

        consistent_count = sum(1 for r in cls_results if r.get('region_consistent', False))
        avg_scores = {k: 0.0 for k in ['optic_disc', 'macula', 'periphery', 'overall']}
        for r in cls_results:
            for k in avg_scores:
                avg_scores[k] += r['region_scores'].get(k, 0.0)
        for k in avg_scores:
            avg_scores[k] = round(avg_scores[k] / len(cls_results), 4)

        dominant_zone = max(
            ['optic_disc', 'macula', 'periphery'],
            key=lambda k: avg_scores[k]
        )

        validation[cls_name] = {
            'n_samples':          len(cls_results),
            'expected_region':    EXPECTED_REGIONS[cls_idx],
            'dominant_zone':      dominant_zone,
            'consistent_samples': consistent_count,
            'consistency_pct':    round(100 * consistent_count / len(cls_results), 1),
            'avg_region_scores':  avg_scores,
        }

        print(f'  {cls_name:15s}: {consistent_count}/{len(cls_results)} consistent '
              f'({validation[cls_name]["consistency_pct"]:.0f}%)  '
              f'dominant={dominant_zone}')

    # Save
    val_path = os.path.join(GRADCAM_DIR, 'heatmap_validation.json')
    with open(val_path, 'w') as f:
        json.dump(validation, f, indent=2)
    print(f'  Validation saved -> {val_path}')

    return validation


# ================================================================
# CLINICAL REPORT
# ================================================================
def generate_clinical_report(results, validation, ood_stats, failed):
    """Generate GRADCAM_REPORT.md with clinical analysis."""
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    n_total    = len(results)
    n_correct  = sum(1 for r in results if r.get('predicted_label') == r.get('true_label'))
    n_ood      = sum(1 for r in results if r.get('ood_flag'))
    n_review   = sum(1 for r in results if r.get('review_flag'))
    avg_conf   = np.mean([r['confidence'] for r in results]) if results else 0.0

    lines = [
        '# RetinaSense v3.0 — Grad-CAM Clinical Report',
        f'',
        f'**Generated**: {now}  ',
        f'**Model**: ViT-Base-Patch16-224 (81.19% test accuracy)  ',
        f'**Pipeline**: Grad-CAM + Mahalanobis OOD + Temperature Scaling + Per-Class Thresholds',
        '',
        '---',
        '',
        '## Executive Summary',
        '',
        f'| Metric | Value |',
        f'|--------|-------|',
        f'| Images processed | {n_total} |',
        f'| Correct predictions | {n_correct}/{n_total} ({100*n_correct/max(n_total,1):.1f}%) |',
        f'| Avg calibrated confidence | {avg_conf:.3f} |',
        f'| OOD flags raised | {n_ood} |',
        f'| Human review flags | {n_review} |',
        f'| Failed images | {len(failed)} |',
        f'| Temperature T | {TEMPERATURE:.4f} |',
        '',
        '---',
        '',
        '## Per-Sample Predictions',
        '',
        '| # | Image | True | Predicted | Confidence | OOD | Review | Attention Region |',
        '|---|-------|------|-----------|-----------|-----|--------|-----------------|',
    ]

    for i, r in enumerate(results):
        true_name = CLASS_NAMES[r['true_label']] if r.get('true_label') is not None else 'Unknown'
        correct_marker = 'OK' if r['predicted_label'] == r.get('true_label') else '**WRONG**'
        lines.append(
            f'| {i+1} | {os.path.basename(r["image_path"])[:25]} '
            f'| {true_name} '
            f'| {r["predicted_class"]} ({correct_marker}) '
            f'| {r["confidence"]:.3f} '
            f'| {"YES" if r["ood_flag"] else "no"} '
            f'| {"YES" if r["review_flag"] else "no"} '
            f'| {r["attention_region"]} |'
        )

    lines += [
        '',
        '---',
        '',
        '## Per-Class Attention Pattern Analysis',
        '',
        '| Disease | Expected Region | Dominant Zone | Consistency |',
        '|---------|----------------|---------------|-------------|',
    ]
    for cls_name, v in validation.items():
        if v.get('n_samples', 0) == 0:
            lines.append(f'| {cls_name} | N/A | N/A | N/A (no samples) |')
        else:
            lines.append(
                f'| {cls_name} | {v["expected_region"]} '
                f'| {v["dominant_zone"]} '
                f'| {v["consistency_pct"]:.0f}% ({v["consistent_samples"]}/{v["n_samples"]}) |'
            )

    lines += [
        '',
        '---',
        '',
        '## OOD Detection Statistics',
        '',
        f'- **Method**: Mahalanobis distance to nearest class centroid (CLS token features)',
        f'- **Threshold percentile**: 97.5th percentile of training-set distances',
        f'- **OOD threshold**: {ood_stats.get("threshold", "N/A")}',
        f'- **Images flagged OOD**: {n_ood}/{n_total}',
        '',
        '### Interpretation',
        '',
        '- Mahalanobis distance measures how far a feature embedding lies from known class distributions',
        '- Low-quality images, extreme artefacts, or off-distribution fundus cameras may trigger OOD flags',
        '- All OOD-flagged images are automatically sent for human review',
        '',
        '---',
        '',
        '## Grad-CAM Heatmap Descriptions',
        '',
        '| Disease | Expected activation | Clinical significance |',
        '|---------|--------------------|-----------------------|',
        '| Normal | Low, uniform | No focal pathology — model attention diffuse |',
        '| Diabetes/DR | Scattered periphery + macula | Microaneurysms, exudates, NV |',
        '| Glaucoma | Optic disc (centre) | Structural disc changes, CDR |',
        '| Cataract | Diffuse lens opacity | Posterior/anterior capsule opacification |',
        '| AMD | Macula / centre-temporal | Drusen, RPE atrophy, CNV |',
        '',
        '---',
        '',
        '## Thresholds Applied',
        '',
        '| Class | Threshold |',
        '|-------|-----------|',
    ]
    for cls_name, thr in zip(CLASS_NAMES, THRESHOLDS):
        lines.append(f'| {cls_name} | {thr:.4f} |')

    lines += [
        '',
        '---',
        '',
        '## Deployment Recommendations',
        '',
        '1. **Confidence gate**: Flag predictions below 0.50 for mandatory ophthalmologist review.',
        '2. **OOD gate**: Any Mahalanobis distance above threshold should trigger QC check on image quality before clinical use.',
        '3. **Grad-CAM review**: Clinicians should inspect heatmaps for cases where model attention does not align with expected anatomy.',
        '4. **Glaucoma caution**: Current dataset imbalance (46 test samples) — consider supplementing ODIR with additional glaucoma images.',
        '5. **Continuous monitoring**: Re-calibrate temperature and thresholds quarterly on production data.',
        '6. **Not for standalone diagnosis**: Grad-CAM is an explainability aid; all predictions require clinical validation.',
        '',
        '---',
        '',
        f'*Report auto-generated by RetinaSense v3.0 Grad-CAM Pipeline | {now}*',
    ]

    report_path = os.path.join(GRADCAM_DIR, 'GRADCAM_REPORT.md')
    with open(report_path, 'w') as f:
        f.write('\n'.join(lines))
    print(f'  Clinical report saved -> {report_path}')
    return report_path


# ================================================================
# MAIN
# ================================================================
def main():
    t_start = time.time()

    # ---- 1. Build Grad-CAM ---
    print('\n[1/6] Initialising ViTGradCAM...')
    gradcam = ViTGradCAM(model)
    print(f'  Method        : Attention Rollout (all 12 transformer blocks)')
    print(f'  Hooks         : {len(gradcam._hooks)} attention hooks registered')
    print(f'  fused_attn disabled for attention weight access')

    # ---- 2. Fit OOD Detector ---
    print('\n[2/6] Fitting OOD detector...')
    ood_path = os.path.join(OUTPUT_DIR, 'ood_detector')
    ood_detector = OODDetector(threshold_percentile=97.5)

    if os.path.exists(ood_path + '.npz'):
        ood_detector.load(ood_path)
    else:
        # Build a small DataLoader from training data to fit OOD detector
        import pandas as pd
        from torch.utils.data import Dataset, DataLoader
        from torchvision import transforms as T

        train_df = pd.read_csv(os.path.join(BASE_DIR, 'data', 'train_split.csv'))

        class SimpleDataset(Dataset):
            def __init__(self, df):
                self.df = df.reset_index(drop=True)
                self.transform = transforms.Compose([
                    transforms.ToPILImage(),
                    transforms.ToTensor(),
                    transforms.Normalize(NORM_MEAN, NORM_STD),
                ])

            def __len__(self):
                return len(self.df)

            def __getitem__(self, idx):
                row = self.df.iloc[idx]
                img_path = str(row['image_path'])
                if not os.path.isabs(img_path):
                    clean = img_path
                    while clean.startswith('./') or clean.startswith('.//'):
                        clean = clean[2:] if clean.startswith('./') else clean[3:]
                    img_path = os.path.join(BASE_DIR, clean)
                dataset = str(row.get('dataset', 'auto'))

                try:
                    img_np, _ = load_and_preprocess(img_path, dataset=dataset)
                    img_tensor = self.transform(img_np)
                except Exception:
                    img_tensor = torch.zeros(3, IMG_SIZE, IMG_SIZE)

                lbl = int(row['disease_label'])
                return img_tensor, torch.tensor(lbl, dtype=torch.long), torch.tensor(0, dtype=torch.long)

        ood_ds     = SimpleDataset(train_df)
        ood_loader = DataLoader(ood_ds, batch_size=32, shuffle=False, num_workers=4)
        ood_detector.fit(model, ood_loader, DEVICE, max_batches=80)
        ood_detector.save(ood_path)

    # ---- 3. Batch Evaluation ---
    print('\n[3/6] Batch evaluation on 20 test images...')
    results, failed = run_batch_evaluation(
        model, gradcam, ood_detector,
        THRESHOLDS, TEMPERATURE, DEVICE,
        n_per_class=4
    )

    # ---- 4. Summary Grid ---
    print('\n[4/6] Generating summary grid...')
    grid_path = save_summary_grid(results)

    # ---- 5. Heatmap Validation ---
    print('\n[5/6] Heatmap validation...')
    validation = validate_heatmaps(results)

    # ---- 6. Clinical Report ---
    print('\n[6/6] Generating clinical report...')
    ood_stats = {'threshold': round(ood_detector.ood_threshold, 4) if ood_detector.is_fitted else 'N/A'}
    report_path = generate_clinical_report(results, validation, ood_stats, failed)

    # ---- Cleanup ---
    gradcam.remove_hooks()

    # ================================================================
    # FINAL SUMMARY
    # ================================================================
    elapsed = time.time() - t_start
    n_total   = len(results)
    n_correct = sum(1 for r in results if r.get('predicted_label') == r.get('true_label'))
    avg_conf  = np.mean([r['confidence'] for r in results]) if results else 0.0
    n_ood     = sum(1 for r in results if r['ood_flag'])
    n_review  = sum(1 for r in results if r['review_flag'])

    print('\n' + '=' * 65)
    print('       RetinaSense v3.0 — GRAD-CAM PIPELINE COMPLETE')
    print('=' * 65)
    print(f'  Images processed     : {n_total}')
    print(f'  Correct predictions  : {n_correct}/{n_total}  ({100*n_correct/max(n_total,1):.1f}%)')
    print(f'  Avg calibrated conf  : {avg_conf:.3f}')
    print(f'  OOD flags            : {n_ood}')
    print(f'  Review flags         : {n_review}')
    print(f'  Failed images        : {len(failed)}')
    print(f'  Elapsed time         : {elapsed:.1f}s')
    print()
    print(f'  Output directory     : {GRADCAM_DIR}')
    output_files = [
        'gradcam_summary_grid.png',
        'heatmap_validation.json',
        'GRADCAM_REPORT.md',
    ] + [os.path.basename(r.get('save_path', '')) for r in results if r.get('save_path')]
    for fname in output_files:
        if fname:
            full = os.path.join(GRADCAM_DIR, fname)
            exists = os.path.exists(full)
            print(f'    {"[OK]" if exists else "[!!]"} {fname}')
    print('=' * 65)

    return results, validation


if __name__ == '__main__':
    main()