integrated_gradients_xai.py

#!/usr/bin/env python3
"""
RetinaSense v3.0 -- Phase 1C: Advanced XAI with Integrated Gradients
=====================================================================
Compares Attention Rollout (existing) vs Integrated Gradients (captum)
on 20 test images (4 per class).

Outputs (all saved to outputs_v3/xai/):
  - comparison_grid.png       : 20-row x 3-column grid [Original | Rollout | IG]
  - ig_individual_01..20.png  : Individual IG heatmaps
  - agreement_heatmap.png     : Spatial correlation matrix between methods
  - agreement_score.json      : Numerical agreement scores per image

Usage:
  python integrated_gradients_xai.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
from matplotlib.colors import Normalize
from PIL import Image
import pandas as pd
from scipy.stats import pearsonr

warnings.filterwarnings('ignore')

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

from captum.attr import IntegratedGradients

# Maximize CPU parallelism
torch.set_num_threads(os.cpu_count() or 4)

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

MODEL_PATH       = os.path.join(OUTPUT_DIR, 'best_model.pth')
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
N_PER_CLASS = 4

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

print('=' * 65)
print('   RetinaSense v3.0 -- Phase 1C: Integrated Gradients XAI')
print('=' * 65)
print(f'  Device   : {DEVICE}')
if torch.cuda.is_available():
    print(f'  GPU      : {torch.cuda.get_device_name(0)}')
print(f'  Output   : {XAI_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]}, '
          f'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')

with open(TEMPERATURE_PATH) as f:
    TEMPERATURE = json.load(f)['temperature']
print(f'  Temperature T = {TEMPERATURE:.4f}')


# ================================================================
# MODEL ARCHITECTURE (mirrors gradcam_v3.py / retinasense_v3.py)
# ================================================================
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

        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)
        f = self.drop(f)
        return self.disease_head(f), self.severity_head(f)


# ================================================================
# DISEASE-LOGITS WRAPPER FOR CAPTUM
# ================================================================
class DiseaseLogitModel(nn.Module):
    """
    Wraps MultiTaskViT so that forward(x) returns only the disease logits.
    Captum's IntegratedGradients requires a model whose forward output
    is either a scalar or a 1-D tensor. We select the target class
    logit inside the IG call via the `target` parameter, so here we
    return the full (B, 5) disease logits.
    """

    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, x):
        disease_logits, _ = self.model(x)
        return disease_logits


# ================================================================
# ATTENTION ROLLOUT (copied from gradcam_v3.py for self-containment)
# ================================================================
class ViTAttentionRollout:
    """
    Attention Rollout for Vision Transformer.
    Traces information flow from patches to CLS token across all layers.
    """

    def __init__(self, model, discard_ratio=0.97):
        self.model = model
        self.discard_ratio = discard_ratio
        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())

    def generate(self, image_tensor, class_idx=None):
        """
        Generate attention rollout heatmap.
        Returns:
            heatmap (224, 224) float32 [0, 1], predicted_label, confidence
        """
        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

        attn_stack = torch.stack(self._attention_maps, dim=0)[:, 0]
        attn_mean = attn_stack.mean(dim=1)

        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)
            attn_mean = torch.where(
                attn_mean >= thresh, attn_mean, torch.zeros_like(attn_mean)
            )

        I = torch.eye(attn_mean.shape[-1]).unsqueeze(0)
        attn_aug = attn_mean + I
        attn_aug = attn_aug / attn_aug.sum(dim=-1, keepdim=True).clamp(min=1e-8)

        rollout = attn_aug[0]
        for l in range(1, len(attn_aug)):
            rollout = rollout @ attn_aug[l]

        cls_attention = rollout[0, 1:]
        spatial = cls_attention.numpy().reshape(14, 14).astype(np.float32)
        spatial = cv2.resize(spatial, (IMG_SIZE, IMG_SIZE),
                             interpolation=cv2.INTER_LINEAR)

        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)

        spatial = np.power(spatial, 0.4)  # gamma stretch

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

    def remove_hooks(self):
        for h in self._hooks:
            h.remove()
        self._hooks = []


# ================================================================
# IMAGE PREPROCESSING (mirrors gradcam_v3.py)
# ================================================================
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 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 with domain-conditional preprocessing.
    Returns:
        img_np   : (224, 224, 3) uint8 preprocessed
        img_orig : (224, 224, 3) uint8 original
    """
    if not os.path.isabs(image_path):
        clean = image_path
        while clean.startswith('./'):
            clean = clean[2:]
        image_path = os.path.join(BASE_DIR, clean)

    if dataset == 'auto':
        if 'aptos' in image_path.lower() or 'gaussian' in image_path.lower():
            dataset = 'APTOS'
        else:
            dataset = 'ODIR'

    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))

    if dataset == 'APTOS':
        img_np = ben_graham(image_path, sz=IMG_SIZE)
    else:
        img_np = clahe_preprocess(image_path, sz=IMG_SIZE)

    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)


# ================================================================
# CIRCULAR FUNDUS MASK
# ================================================================
def create_fundus_mask(h=IMG_SIZE, w=IMG_SIZE):
    """
    Create a soft circular mask matching the fundus region.
    Uses a smooth Gaussian-blurred edge to avoid hard boundaries.
    Returns float32 mask [0, 1] of shape (h, w).
    """
    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)
    return mask


# ================================================================
# INTEGRATED GRADIENTS COMPUTATION
# ================================================================
def compute_ig_attribution(ig_model, ig_method, img_tensor, target_class,
                           n_steps=50, internal_batch_size=10, sigma=10):
    """
    Compute Integrated Gradients attribution for a single image.

    Uses a Gaussian-blurred baseline (sigma=10) which is more appropriate
    for fundus images than a black baseline (since the background is already dark).

    Args:
        ig_model          : DiseaseLogitModel wrapper
        ig_method         : captum IntegratedGradients instance
        img_tensor        : (1, 3, 224, 224) normalised tensor on DEVICE
        target_class      : int, disease class to explain
        n_steps           : number of interpolation steps
        internal_batch_size : batch size for internal IG computation
        sigma             : Gaussian blur sigma for baseline

    Returns:
        attribution : (224, 224) float32 numpy array, normalised [0, 1]
    """
    # Create blurred baseline in pixel space, then normalise
    # First undo normalisation to get pixel-space tensor
    mean_t = torch.tensor(NORM_MEAN, device=DEVICE).view(1, 3, 1, 1)
    std_t  = torch.tensor(NORM_STD, device=DEVICE).view(1, 3, 1, 1)

    # Build the blurred baseline from the input tensor
    # Denormalise -> blur -> renormalise
    img_denorm = img_tensor * std_t + mean_t  # approx [0, 1] range
    img_np_for_blur = (img_denorm[0].permute(1, 2, 0).cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
    blurred_np = cv2.GaussianBlur(img_np_for_blur, (0, 0), sigma)
    # Convert blurred back to tensor with normalisation
    blurred_tensor = transforms.Compose([
        transforms.ToPILImage(),
        transforms.ToTensor(),
        transforms.Normalize(NORM_MEAN, NORM_STD),
    ])(blurred_np).unsqueeze(0).to(DEVICE)

    # Compute Integrated Gradients
    img_tensor.requires_grad_(True)
    attributions = ig_method.attribute(
        img_tensor,
        baselines=blurred_tensor,
        target=target_class,
        n_steps=n_steps,
        internal_batch_size=internal_batch_size,
    )

    # Aggregate across channels: take L2 norm across RGB for spatial map
    # shape: (1, 3, 224, 224) -> (224, 224)
    attr_np = attributions[0].detach().cpu().numpy()  # (3, 224, 224)
    # Use absolute values and sum over channels for a positive attribution map
    attr_spatial = np.sqrt(np.sum(attr_np ** 2, axis=0))  # (224, 224)

    # Normalise to [0, 1]
    a_min, a_max = attr_spatial.min(), attr_spatial.max()
    if a_max - a_min > 1e-8:
        attr_spatial = (attr_spatial - a_min) / (a_max - a_min)
    else:
        attr_spatial = np.zeros_like(attr_spatial)

    return attr_spatial.astype(np.float32)


# ================================================================
# OVERLAY FUNCTION
# ================================================================
def overlay_heatmap(original_np, heatmap, alpha=0.6, cmap_name='inferno'):
    """
    Blend heatmap onto original image with circular fundus mask.

    Args:
        original_np : (224, 224, 3) uint8 RGB
        heatmap     : (224, 224) float32 [0, 1]
        alpha       : heatmap blending opacity
        cmap_name   : matplotlib colormap name

    Returns:
        blended : (224, 224, 3) uint8 RGB
    """
    # Apply colormap
    cmap = plt.get_cmap(cmap_name)
    colored = cmap(heatmap)[:, :, :3]  # (224, 224, 3) float [0, 1]
    colored_uint8 = (colored * 255).astype(np.uint8)

    # Get fundus mask
    mask = create_fundus_mask(heatmap.shape[0], heatmap.shape[1])

    # Blend inside the fundus region only
    orig = original_np.astype(np.float32)
    cmap_f = colored_uint8.astype(np.float32)
    blended = orig.copy()
    for c in range(3):
        blended[:, :, c] = (
            orig[:, :, c] * (1 - alpha * mask)
            + cmap_f[:, :, c] * (alpha * mask)
        )
    return np.clip(blended, 0, 255).astype(np.uint8)


# ================================================================
# SELECT TEST IMAGES (same logic as gradcam_v3.py)
# ================================================================
def select_test_images(n_per_class=N_PER_CLASS):
    """Select n_per_class images per disease class from test split."""
    df = pd.read_csv(TEST_CSV)
    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('source', 'auto')),
            })
    print(f'  Selected {len(samples)} test images '
          f'({n_per_class} per class x {NUM_CLASSES} classes)')
    return samples


# ================================================================
# COMPUTE AGREEMENT METRICS
# ================================================================
def compute_agreement(rollout_map, ig_map, fundus_mask):
    """
    Compute spatial agreement between Attention Rollout and IG heatmaps.

    Metrics:
        - Pearson correlation (within fundus mask)
        - Intersection over Union (IoU) of top-20% activated regions
        - Cosine similarity of flattened masked vectors

    Returns dict of scores.
    """
    # Flatten inside mask
    mask_bool = fundus_mask > 0.5
    r_flat = rollout_map[mask_bool]
    i_flat = ig_map[mask_bool]

    # Pearson correlation
    if len(r_flat) > 2 and r_flat.std() > 1e-8 and i_flat.std() > 1e-8:
        pearson_r, pearson_p = pearsonr(r_flat, i_flat)
    else:
        pearson_r, pearson_p = 0.0, 1.0

    # IoU of top-20% regions
    r_thresh = np.percentile(r_flat, 80)
    i_thresh = np.percentile(i_flat, 80)
    r_top = rollout_map > r_thresh
    i_top = ig_map > i_thresh
    intersection = np.logical_and(r_top, i_top).sum()
    union = np.logical_or(r_top, i_top).sum()
    iou = float(intersection / max(union, 1))

    # Cosine similarity
    r_norm = np.linalg.norm(r_flat)
    i_norm = np.linalg.norm(i_flat)
    if r_norm > 1e-8 and i_norm > 1e-8:
        cosine = float(np.dot(r_flat, i_flat) / (r_norm * i_norm))
    else:
        cosine = 0.0

    return {
        'pearson_r':     round(float(pearson_r), 4),
        'pearson_p':     round(float(pearson_p), 6),
        'iou_top20':     round(iou, 4),
        'cosine_sim':    round(cosine, 4),
    }


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

    # ---- 1. Load model ----
    print('\n[1/6] Loading 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}')

    # ---- 2. Select images ----
    print('\n[2/6] Selecting test images...')
    samples = select_test_images()

    # ---- 3. Preprocess all images + run Attention Rollout ----
    # IMPORTANT: Run rollout FIRST then remove hooks BEFORE IG.
    # This prevents rollout hooks from firing during IG's many
    # forward passes (50 steps), which would be very slow on CPU.
    print('\n[3/6] Running Attention Rollout on all images...')
    rollout = ViTAttentionRollout(model, discard_ratio=0.97)
    print(f'  Attention Rollout: {len(rollout._hooks)} hooks registered')

    preprocessed = []  # store (img_np, img_orig, img_tensor) per sample
    rollout_results = []  # store (heatmap, pred_label, confidence) per sample

    for idx, sample in enumerate(samples):
        img_path   = sample['image_path']
        true_label = sample['true_label']
        dataset    = sample['dataset']
        basename   = os.path.basename(img_path)

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

        try:
            img_np, img_orig = load_and_preprocess(img_path, dataset=dataset)
            img_tensor = preprocess_to_tensor(img_np).to(DEVICE)
            preprocessed.append((img_np, img_orig, img_tensor))

            heatmap, pred_label, pred_conf = rollout.generate(img_tensor)
            rollout_results.append((heatmap, pred_label, pred_conf))
            print(f'-> pred={CLASS_NAMES[pred_label]:15s}  conf={pred_conf:.3f}')

        except Exception as e:
            print(f'FAILED: {e}')
            preprocessed.append(None)
            rollout_results.append(None)

    # Remove rollout hooks BEFORE running IG to avoid extra computation
    rollout.remove_hooks()
    # Re-enable fused attention for faster forward passes during IG
    for block in model.backbone.blocks:
        block.attn.fused_attn = True
    print('  Rollout hooks removed. fused_attn re-enabled for IG speed.')

    # ---- 4. Run Integrated Gradients (no rollout hooks active) ----
    print('\n[4/6] Computing Integrated Gradients attributions...')
    disease_model = DiseaseLogitModel(model)
    disease_model.eval()
    ig_method = IntegratedGradients(disease_model)
    print(f'  Baseline: Gaussian blur (sigma=10), n_steps=50, '
          f'internal_batch_size=25')

    all_results = []
    fundus_mask = create_fundus_mask()

    for idx, sample in enumerate(samples):
        if preprocessed[idx] is None or rollout_results[idx] is None:
            continue

        img_path   = sample['image_path']
        true_label = sample['true_label']
        basename   = os.path.basename(img_path)
        img_np, img_orig, img_tensor = preprocessed[idx]
        rollout_heatmap, pred_label, pred_conf = rollout_results[idx]

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

        try:
            # Fresh tensor (IG needs requires_grad)
            ig_input = img_tensor.clone().detach().to(DEVICE)

            ig_heatmap = compute_ig_attribution(
                disease_model, ig_method, ig_input,
                target_class=pred_label,
                n_steps=50,
                internal_batch_size=25,
                sigma=10,
            )

            # Agreement
            agreement = compute_agreement(rollout_heatmap, ig_heatmap,
                                          fundus_mask)

            print(f'-> pearson={agreement["pearson_r"]:.3f}  '
                  f'IoU={agreement["iou_top20"]:.3f}')

            all_results.append({
                'idx':              idx,
                'image_path':       img_path,
                'basename':         basename,
                'true_label':       true_label,
                'pred_label':       pred_label,
                'pred_class':       CLASS_NAMES[pred_label],
                'confidence':       round(pred_conf, 4),
                'img_orig':         img_orig,
                'rollout_heatmap':  rollout_heatmap,
                'ig_heatmap':       ig_heatmap,
                'agreement':        agreement,
            })

        except Exception as e:
            print(f'FAILED: {e}')
            import traceback
            traceback.print_exc()
            continue

    n_success = len(all_results)
    print(f'\n  Completed: {n_success}/{len(samples)} images')

    if n_success == 0:
        print('ERROR: No images processed successfully. Exiting.')
        sys.exit(1)

    # ---- 5. Generate visualisations ----
    print('\n[5/6] Generating visualisations...')

    # 5a. Individual IG heatmaps
    print('  Saving individual IG heatmaps...')
    for r in all_results:
        fig, axes = plt.subplots(1, 3, figsize=(13, 4.5))

        # Original
        axes[0].imshow(r['img_orig'])
        axes[0].set_title(f'Original\nTrue: {CLASS_NAMES[r["true_label"]]}',
                          fontsize=10, fontweight='bold')
        axes[0].axis('off')

        # IG heatmap (raw)
        im = axes[1].imshow(r['ig_heatmap'], cmap='inferno', vmin=0, vmax=1)
        axes[1].set_title('Integrated Gradients\n(attribution magnitude)',
                          fontsize=10)
        axes[1].axis('off')

        # IG overlay on original
        ig_overlay = overlay_heatmap(r['img_orig'], r['ig_heatmap'],
                                     alpha=0.6, cmap_name='inferno')
        axes[2].imshow(ig_overlay)
        correct = r['pred_label'] == r['true_label']
        status = 'OK' if correct else 'WRONG'
        axes[2].set_title(
            f'IG Overlay\nPred: {r["pred_class"]} ({r["confidence"]:.2f}) [{status}]',
            fontsize=10,
            color='green' if correct else 'red',
            fontweight='bold')
        axes[2].axis('off')

        plt.tight_layout()
        save_path = os.path.join(XAI_DIR,
                                 f'ig_individual_{r["idx"]+1:02d}.png')
        fig.savefig(save_path, dpi=150, bbox_inches='tight',
                    facecolor='white')
        plt.close(fig)

    print(f'    Saved {n_success} individual IG heatmaps')

    # 5b. Comparison grid: 20 rows x 3 columns
    print('  Generating comparison grid...')
    n_rows = n_success
    fig, axes = plt.subplots(n_rows, 3, figsize=(14, 4.2 * n_rows))

    # Handle single-row case
    if n_rows == 1:
        axes = axes[np.newaxis, :]

    # Column headers
    col_titles = ['Original Image', 'Attention Rollout', 'Integrated Gradients']

    for i, r in enumerate(all_results):
        true_name = CLASS_NAMES[r['true_label']]
        pred_name = r['pred_class']
        correct = r['pred_label'] == r['true_label']
        status = 'OK' if correct else 'WRONG'

        # Column 0: Original
        axes[i, 0].imshow(r['img_orig'])
        axes[i, 0].set_ylabel(f'#{r["idx"]+1}\nTrue: {true_name}',
                               fontsize=9, fontweight='bold', rotation=0,
                               labelpad=70, va='center')
        axes[i, 0].set_xticks([])
        axes[i, 0].set_yticks([])
        if i == 0:
            axes[i, 0].set_title(col_titles[0], fontsize=12, fontweight='bold',
                                 pad=10)

        # Column 1: Attention Rollout overlay
        rollout_overlay = overlay_heatmap(r['img_orig'], r['rollout_heatmap'],
                                          alpha=0.6, cmap_name='inferno')
        axes[i, 1].imshow(rollout_overlay)
        axes[i, 1].axis('off')
        if i == 0:
            axes[i, 1].set_title(col_titles[1], fontsize=12, fontweight='bold',
                                 pad=10)

        # Column 2: IG overlay
        ig_overlay = overlay_heatmap(r['img_orig'], r['ig_heatmap'],
                                     alpha=0.6, cmap_name='inferno')
        axes[i, 2].imshow(ig_overlay)
        color = 'green' if correct else 'red'
        axes[i, 2].set_xlabel(
            f'Pred: {pred_name} ({r["confidence"]:.2f}) [{status}] | '
            f'Pearson r={r["agreement"]["pearson_r"]:.2f}',
            fontsize=8, color=color, fontweight='bold')
        axes[i, 2].set_xticks([])
        axes[i, 2].set_yticks([])
        if i == 0:
            axes[i, 2].set_title(col_titles[2], fontsize=12, fontweight='bold',
                                 pad=10)

    plt.suptitle('RetinaSense v3.0 -- Attention Rollout vs Integrated Gradients',
                 fontsize=14, fontweight='bold', y=1.001)
    plt.tight_layout()
    grid_path = os.path.join(XAI_DIR, 'comparison_grid.png')
    fig.savefig(grid_path, dpi=120, bbox_inches='tight', facecolor='white')
    plt.close(fig)
    print(f'    Saved: {grid_path}')

    # 5c. Agreement heatmap (matrix showing per-image spatial correlation)
    print('  Generating agreement heatmap...')

    # Build per-image metrics matrix
    image_labels = [
        f'#{r["idx"]+1} {CLASS_NAMES[r["true_label"]][:6]}'
        for r in all_results
    ]
    metric_names = ['Pearson r', 'IoU (top 20%)', 'Cosine Sim']
    agreement_matrix = np.zeros((n_success, 3))
    for i, r in enumerate(all_results):
        agreement_matrix[i, 0] = r['agreement']['pearson_r']
        agreement_matrix[i, 1] = r['agreement']['iou_top20']
        agreement_matrix[i, 2] = r['agreement']['cosine_sim']

    fig, ax = plt.subplots(figsize=(7, max(8, n_success * 0.45)))
    im = ax.imshow(agreement_matrix, cmap='RdYlGn', aspect='auto',
                   vmin=-0.2, vmax=1.0)

    ax.set_xticks(range(3))
    ax.set_xticklabels(metric_names, fontsize=10, fontweight='bold')
    ax.set_yticks(range(n_success))
    ax.set_yticklabels(image_labels, fontsize=8)

    # Annotate cells
    for i in range(n_success):
        for j in range(3):
            val = agreement_matrix[i, j]
            color = 'white' if val < 0.3 else 'black'
            ax.text(j, i, f'{val:.2f}', ha='center', va='center',
                    fontsize=8, color=color, fontweight='bold')

    ax.set_title('Rollout vs IG Agreement Scores per Image',
                 fontsize=12, fontweight='bold', pad=12)
    plt.colorbar(im, ax=ax, shrink=0.6, label='Score')
    plt.tight_layout()

    heatmap_path = os.path.join(XAI_DIR, 'agreement_heatmap.png')
    fig.savefig(heatmap_path, dpi=150, bbox_inches='tight', facecolor='white')
    plt.close(fig)
    print(f'    Saved: {heatmap_path}')

    # ---- 6. Save agreement scores JSON ----
    print('\n[6/6] Saving agreement scores...')

    scores_output = {
        'summary': {
            'n_images':          n_success,
            'mean_pearson_r':    round(float(agreement_matrix[:, 0].mean()), 4),
            'mean_iou_top20':    round(float(agreement_matrix[:, 1].mean()), 4),
            'mean_cosine_sim':   round(float(agreement_matrix[:, 2].mean()), 4),
            'std_pearson_r':     round(float(agreement_matrix[:, 0].std()), 4),
            'std_iou_top20':     round(float(agreement_matrix[:, 1].std()), 4),
            'std_cosine_sim':    round(float(agreement_matrix[:, 2].std()), 4),
        },
        'per_image': [],
    }
    for r in all_results:
        scores_output['per_image'].append({
            'image':       r['basename'],
            'true_label':  r['true_label'],
            'true_class':  CLASS_NAMES[r['true_label']],
            'pred_label':  r['pred_label'],
            'pred_class':  r['pred_class'],
            'confidence':  r['confidence'],
            'agreement':   r['agreement'],
        })

    # Per-class summary
    per_class = {}
    for cls_idx in range(NUM_CLASSES):
        cls_results = [r for r in all_results if r['true_label'] == cls_idx]
        if cls_results:
            pearson_vals = [r['agreement']['pearson_r'] for r in cls_results]
            iou_vals     = [r['agreement']['iou_top20'] for r in cls_results]
            cosine_vals  = [r['agreement']['cosine_sim'] for r in cls_results]
            per_class[CLASS_NAMES[cls_idx]] = {
                'n_images':        len(cls_results),
                'mean_pearson_r':  round(float(np.mean(pearson_vals)), 4),
                'mean_iou_top20':  round(float(np.mean(iou_vals)), 4),
                'mean_cosine_sim': round(float(np.mean(cosine_vals)), 4),
            }
    scores_output['per_class'] = per_class

    json_path = os.path.join(XAI_DIR, 'agreement_score.json')
    with open(json_path, 'w') as f:
        json.dump(scores_output, f, indent=2)
    print(f'  Saved: {json_path}')

    # ---- Summary ----
    elapsed = time.time() - t_start
    n_correct = sum(1 for r in all_results
                    if r['pred_label'] == r['true_label'])

    print('\n' + '=' * 65)
    print('   PHASE 1C COMPLETE: Integrated Gradients XAI')
    print('=' * 65)
    print(f'  Images processed  : {n_success}/{len(samples)}')
    print(f'  Correct preds     : {n_correct}/{n_success} '
          f'({100*n_correct/max(n_success,1):.1f}%)')
    print(f'  Mean Pearson r    : {scores_output["summary"]["mean_pearson_r"]:.4f}')
    print(f'  Mean IoU (top 20%): {scores_output["summary"]["mean_iou_top20"]:.4f}')
    print(f'  Mean Cosine Sim   : {scores_output["summary"]["mean_cosine_sim"]:.4f}')
    print(f'  Time elapsed      : {elapsed:.1f}s')
    print(f'  Outputs in        : {XAI_DIR}')
    print('=' * 65)

    # List outputs
    print('\n  Output files:')
    for fname in sorted(os.listdir(XAI_DIR)):
        fpath = os.path.join(XAI_DIR, fname)
        size_kb = os.path.getsize(fpath) / 1024
        print(f'    {fname:40s}  {size_kb:8.1f} KB')


if __name__ == '__main__':
    main()