Session 6 — Variational Inference: Implementation 🏥¶

!pip install -q monai torch-uncertainty lightning

import os, random, glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from PIL import Image
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim as optim

from scipy import stats
from statsmodels.stats.proportion import proportions_ztest
from sklearn.metrics import roc_auc_score, accuracy_score, confusion_matrix, classification_report
from sklearn.mixture import GaussianMixture

# MONAI
from monai.transforms import (
    Compose, LoadImaged, EnsureChannelFirstd, RepeatChanneld,
    ScaleIntensityd, Resized, RandFlipd, RandRotated, EnsureTyped,
)
from monai.data import Dataset, DataLoader
from monai.networks.nets import DenseNet121

# Torch-Uncertainty
from torch_uncertainty.layers.bayesian import BayesLinear
from torch_uncertainty.losses import ELBOLoss

# Reproducibility
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Device: {device}')

DATA_DIR = "/kaggle/input/datasets/paultimothymooney/chest-xray-pneumonia/chest_xray"

def is_grayscale(path):
    try:
        return Image.open(path).mode in ['L', 'I;16', 'I']
    except:
        return False

def make_split_list_filtered(split_dir):
    normal_files = sorted(glob.glob(os.path.join(split_dir, 'NORMAL', '*')))
    pneu_files   = sorted(glob.glob(os.path.join(split_dir, 'PNEUMONIA', '*')))
    normal_files = [f for f in normal_files if is_grayscale(f)]
    pneu_files   = [f for f in pneu_files   if is_grayscale(f)]
    return ([{'image': f, 'label': 0} for f in normal_files] +
            [{'image': f, 'label': 1} for f in pneu_files])

def subsample_with_ratio(files, n_total, normal_pneu_ratio, seed=42):

    rng    = np.random.default_rng(seed)
    normal = [d for d in files if int(d['label']) == 0]
    pneu   = [d for d in files if int(d['label']) == 1]
    rng.shuffle(normal)
    rng.shuffle(pneu)

    n_pneu   = int(round(n_total / (1 + normal_pneu_ratio)))
    n_normal = n_total - n_pneu

    n_normal = min(n_normal, len(normal))
    n_pneu   = min(n_pneu,   len(pneu))

    subset = normal[:n_normal] + pneu[:n_pneu]
    rng.shuffle(subset)

    actual = n_normal / n_pneu if n_pneu > 0 else float('inf')
    print(f'  NORMAL={n_normal}, PNEUMONIA={n_pneu}, '
          f'total={len(subset)}')
    return subset

train_files_all = make_split_list_filtered(os.path.join(DATA_DIR, 'train'))
val_files_all   = make_split_list_filtered(os.path.join(DATA_DIR, 'val'))
test_files_all  = make_split_list_filtered(os.path.join(DATA_DIR, 'test'))

print('Train subset (normal:pneu = 2.7:1):')
train_files = subsample_with_ratio(train_files_all, n_total=1000,
                                    normal_pneu_ratio=2.7, seed=SEED)
print('Test subset (normal:pneu = 1.7:1):')
test_files  = subsample_with_ratio(test_files_all,  n_total=100,
                                    normal_pneu_ratio=1.7, seed=SEED)
print(f'Val (full): {len(val_files_all)} samples')

IMG_SIZE   = 224
BATCH_SIZE = 32

train_tfms = Compose([
    LoadImaged(keys=['image']),
    EnsureChannelFirstd(keys=['image']),
    RepeatChanneld(keys=['image'], repeats=3),   # grayscale [1,H,W] -> [3,H,W]
    ScaleIntensityd(keys=['image']),
    Resized(keys=['image'], spatial_size=(IMG_SIZE, IMG_SIZE)),
    RandFlipd(keys=['image'], prob=0.5, spatial_axis=1),
    RandRotated(keys=['image'], prob=0.3, range_x=0.12),
    EnsureTyped(keys=['image', 'label']),
])

eval_tfms = Compose([
    LoadImaged(keys=['image']),
    EnsureChannelFirstd(keys=['image']),
    RepeatChanneld(keys=['image'], repeats=3),
    ScaleIntensityd(keys=['image']),
    Resized(keys=['image'], spatial_size=(IMG_SIZE, IMG_SIZE)),
    EnsureTyped(keys=['image', 'label']),
])

train_ds = Dataset(data=train_files,   transform=train_tfms)
val_ds   = Dataset(data=val_files_all, transform=eval_tfms)
test_ds  = Dataset(data=test_files,    transform=eval_tfms)

train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE,
                           shuffle=True,  num_workers=2, pin_memory=True)
val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE,
                           shuffle=False, num_workers=2, pin_memory=True)
test_loader  = DataLoader(test_ds,  batch_size=BATCH_SIZE,
                           shuffle=False, num_workers=2, pin_memory=True)

print(f'Train batches : {len(train_loader)}')
print(f'Val batches   : {len(val_loader)}')
print(f'Test batches  : {len(test_loader)}')

# Visualizing one case per class

def get_one_per_class(ds):
    normal_sample = None
    pneumonia_sample = None

    for i in range(len(ds)):
        sample = ds[i]
        y = int(sample["label"])
        if y == 0 and normal_sample is None:
            normal_sample = sample
        elif y == 1 and pneumonia_sample is None:
            pneumonia_sample = sample
        if normal_sample is not None and pneumonia_sample is not None:
            break
    return normal_sample, pneumonia_sample


def to_rgb_for_display(x):
    # x: [C,H,W]
    if x.shape[0] > 3:
        x = x[:3]
    elif x.shape[0] == 1:
        x = x.repeat(3, 1, 1)

    img = x.permute(2, 1, 0).cpu().numpy()
    img = np.clip(img, 0, 1)
    return img


def show_normal_and_pneumonia(ds):
    normal_sample, pneumonia_sample = get_one_per_class(ds)

    plt.figure(figsize=(10, 5))
    for i, sample in enumerate([normal_sample, pneumonia_sample], 1):
        x = sample["image"]
        y = int(sample["label"])
        plt.subplot(1, 2, i)
        plt.imshow(to_rgb_for_display(x))
        plt.title("NORMAL" if y == 0 else "PNEUMONIA")
        plt.axis("off")
    plt.tight_layout()
    plt.show()

show_normal_and_pneumonia(train_ds)

class BayesianDenseNet121(nn.Module):
    """
    MONAI DenseNet121 with a Torch-Uncertainty BayesLinear head.

    The convolutional backbone extracts features deterministically
    using ImageNet pretrained weights.

    The final layer is Bayesian: it learns a distribution over
    weights (mean + variance per weight) instead of a single point.
    This is where Variational Inference lives in the model.
    """

    def __init__(self, num_classes=2, freeze_backbone=True):
        super().__init__()

        # MONAI DenseNet121 with ImageNet pretrained weights
        # spatial_dims=2 for 2D chest X-rays
        # in_channels=3  because RepeatChanneld gives us 3-channel input
        backbone = DenseNet121(
            spatial_dims=2,
            in_channels=3,
            out_channels=num_classes,
            pretrained=True,
        )

        # Keep the feature extraction part
        self.features     = backbone.features
        self.pool         = nn.AdaptiveAvgPool2d((1, 1))

        # Replace the final linear layer with a Bayesian one
        # DenseNet121 final feature dim = 1024
        self.bayes_head = BayesLinear(1024, num_classes)

        # Freeze all backbone parameters — only train the Bayesian head
        if freeze_backbone:
            for param in self.features.parameters():
                param.requires_grad = False

    def forward(self, x):
        x = self.features(x)             # [B, 1024, H', W'] — deterministic
        x = torch.relu(x)               # DenseNet convention: relu after features
        x = self.pool(x)                # [B, 1024, 1, 1]
        x = torch.flatten(x, 1)         # [B, 1024]
        x = self.bayes_head(x)          # [B, 2] — stochastic sampling here
        return x


model = BayesianDenseNet121(num_classes=2, freeze_backbone=True).to(device)

total_params = sum(p.numel() for p in model.parameters())
trainable    = sum(p.numel() for p in model.parameters() if p.requires_grad)
frozen       = total_params - trainable

print(f'Total parameters    : {total_params:,}')
print(f'Frozen (backbone)   : {frozen:,}  <- ImageNet pretrained, fixed')
print(f'Trainable (BayesHead): {trainable:,} <- learns mean + variance per weight')

N_TRAIN    = len(train_ds)
NUM_EPOCHS = 100
T_TRAIN    = 3
T_SAMPLES  = 20

EARLY_STOP_PATIENCE = 10
PLATEAU_PATIENCE    = 3
PLATEAU_FACTOR      = 0.5
MIN_LR              = 1e-6
BEST_PATH           = "/kaggle/working/best_model.pth"

loss_fn = ELBOLoss(
    model       = model,
    inner_loss  = nn.CrossEntropyLoss(),
    kl_weight   = 1 / N_TRAIN,
    num_samples = T_TRAIN,
)

# Only optimize Bayesian head
optimizer = optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3,
)

# Plateau scheduler
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=PLATEAU_FACTOR,
    patience=PLATEAU_PATIENCE,
    threshold=1e-3,
    threshold_mode='rel',
    min_lr=MIN_LR
)

print(f'Training samples     : {N_TRAIN}')
print(f'KL weight            : {1/N_TRAIN:.6f}')
print(f'MC samples (train)   : {T_TRAIN}')
print(f'MC samples (infer)   : {T_SAMPLES}')
print(f'Epochs               : {NUM_EPOCHS}')

history = {
    'train_loss': [],
    'val_loss': [],
    'val_acc': [],
    'val_auc': []
}

best_val_loss = float('inf')
early_stop_counter = 0

for epoch in range(1, NUM_EPOCHS + 1):

    # ── Train ─────────────────────────────────────────────────────
    model.train()
    epoch_loss = 0.0

    for batch in tqdm(train_loader,
                      desc=f'Epoch {epoch:02d}/{NUM_EPOCHS}', leave=False):
        images = batch['image'].to(device, dtype=torch.float32)
        labels = batch['label'].to(device, dtype=torch.long)

        optimizer.zero_grad()
        loss = loss_fn(images, labels)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(train_loader)

    # ── Validate ──────────────────────────────────────
    model.eval()
    val_probs, val_preds, val_labels = [], [], []
    val_loss_total = 0.0

    with torch.no_grad():
        for batch in val_loader:
            images = batch['image'].to(device, dtype=torch.float32)
            labels_tensor = batch['label'].to(device, dtype=torch.long)
            labels_np = labels_tensor.cpu().numpy()

            # ELBO validation loss
            loss = loss_fn(images, labels_tensor)
            val_loss_total += loss.item()

            # predictions
            logits = model(images)
            probs  = torch.softmax(logits, dim=-1).cpu().numpy()

            val_probs.extend(probs[:, 1].tolist())
            val_preds.extend(probs.argmax(axis=1).tolist())
            val_labels.extend(labels_np.tolist())

    val_loss = val_loss_total / len(val_loader)
    val_acc  = accuracy_score(val_labels, val_preds)
    val_auc  = roc_auc_score(val_labels, val_probs)

    # ── Scheduler step ─────────────────────────────
    scheduler.step(val_loss)
    current_lr = optimizer.param_groups[0]['lr']

    # ── Early stopping + best model ────────────────────────────────
    if val_loss < best_val_loss - 1e-4:
        best_val_loss = val_loss
        early_stop_counter = 0

        torch.save(model.state_dict(), BEST_PATH)
        print(f"✅ Saved best model (val_loss={val_loss:.4f})")

    else:
        early_stop_counter += 1

    if early_stop_counter >= EARLY_STOP_PATIENCE:
        print(f"\n⛔ Early stopping triggered at epoch {epoch}")
        break

    # ── Log history ───────────────────────────────────────────────
    history['train_loss'].append(avg_loss)
    history['val_loss'].append(val_loss)
    history['val_acc'].append(val_acc)
    history['val_auc'].append(val_auc)

    print(f'Epoch {epoch:02d}/{NUM_EPOCHS} | '
          f'Train ELBO: {avg_loss:.4f} | '
          f'Val ELBO: {val_loss:.4f} | '
          f'Val Acc: {val_acc:.3f} | '
          f'Val AUC: {val_auc:.3f} | '
          f'LR: {current_lr:.6f}')

print('\nTraining complete ✅')

# Training curves
fig, axes = plt.subplots(1, 2, figsize=(18, 4))
fig.suptitle('Variational Inference — Training curves', fontsize=13)

epochs = range(1, len(history['train_loss']) + 1)

plot_config = [
    ('train_loss', '#D85A30', 'Train ELBO Loss'),
    ('val_loss',   '#7F77DD', 'Val ELBO Loss')
]

for ax, (key, color, title) in zip(axes, plot_config):
    ax.plot(epochs, history[key], color=color, linewidth=2)
    ax.set_title(title)
    ax.set_xlabel('Epoch')
    ax.spines[['top', 'right']].set_visible(False)

plt.tight_layout()
plt.show()

# Loading the best model

model = BayesianDenseNet121(
    num_classes=2,
    freeze_backbone=True
).to(device)

# I renamed the best model to VI_best_model.pth and added it to inputs

BEST_PATH = "/kaggle/input/models/benyamingheiji/vi-best-model-pth/pytorch/default/1/VI_best_model.pth"

model.load_state_dict(torch.load(BEST_PATH, map_location=device))

model.eval()

with torch.no_grad():
    batch = next(iter(test_loader))
    images = batch['image'].to(device)

    logits = model(images)
    probs = torch.softmax(logits, dim=-1)

print("Loaded model works. Output shape:", probs.shape)

def stochastic_inference_vi(model, loader, T=20, device=device):
    """
    VI inference: T forward passes, each sampling different weights from q_φ(w).

    model.eval() keeps batchnorm deterministic.
    BayesLinear still samples weights each forward pass — this is correct VI.

    Returns results DataFrame and raw image tensors.
    """
    model.eval()

    all_mean_preds, all_entropy, all_variance = [], [], []
    all_epistemic, all_aleatoric              = [], []
    all_labels, all_paths, all_images         = [], [], []

    with torch.no_grad():
        for batch in tqdm(loader, desc='VI stochastic inference'):
            images = batch['image'].to(device, dtype=torch.float32)
            labels = batch['label'].cpu().numpy()
            paths  = (batch.get('image_meta_dict', {}).get('filename_or_obj')
                      or [''] * len(labels))

            # T forward passes — each samples different weights from q_φ
            sample_probs = []
            for _ in range(T):
                probs = torch.softmax(model(images), dim=-1).cpu().numpy()
                sample_probs.append(probs)

            sample_probs = np.stack(sample_probs, axis=0)  # [T, B, 2]
            mean_probs   = sample_probs.mean(axis=0)       # [B, 2]

            # Total uncertainty: H of the mean prediction
            total_H = -np.sum(
                mean_probs * np.log(mean_probs + 1e-8), axis=-1
            )  # [B]

            # Aleatoric: expected entropy across T weight samples
            per_H       = -np.sum(
                sample_probs * np.log(sample_probs + 1e-8), axis=-1
            )  # [T, B]
            aleatoric_H = per_H.mean(axis=0)   # [B]

            # Epistemic: mutual information = total - aleatoric
            epistemic_MI = total_H - aleatoric_H   # [B]

            # Variance of P(pneumonia) across T
            variance = sample_probs[:, :, 1].var(axis=0)   # [B]

            all_mean_preds.extend(mean_probs[:, 1].tolist())
            all_entropy.extend(total_H.tolist())
            all_epistemic.extend(epistemic_MI.tolist())
            all_aleatoric.extend(aleatoric_H.tolist())
            all_variance.extend(variance.tolist())
            all_labels.extend(labels.tolist())
            all_paths.extend(list(paths))
            all_images.extend(images.cpu().unbind(0))

    preds   = (np.array(all_mean_preds) >= 0.5).astype(int)
    correct = (preds == np.array(all_labels)).astype(int)

    results = pd.DataFrame({
        'image_path'        : all_paths,
        'true_label'        : all_labels,
        'mean_pred'         : all_mean_preds,
        'pred_label'        : preds.tolist(),
        'predictive_entropy': all_entropy,
        'epistemic_mi'      : all_epistemic,
        'aleatoric_h'       : all_aleatoric,
        'variance'          : all_variance,
        'correct'           : correct.tolist(),
    })
    return results, all_images


results, test_images = stochastic_inference_vi(
    model, test_loader, T=T_SAMPLES, device=device
)

print(f'Test samples  : {len(results)}')
print(f'Accuracy      : {results["correct"].mean():.3f}')
print(f'AUC-ROC       : {roc_auc_score(results["true_label"], results["mean_pred"]):.3f}')

# Uncertainty decomposition

print('Uncertainty decomposition (mean across test set):')
print(f'  Total entropy  : {results["predictive_entropy"].mean():.4f}')
print(f'  Epistemic (MI) : {results["epistemic_mi"].mean():.4f}')
print(f'  Aleatoric (H)  : {results["aleatoric_h"].mean():.4f}')

# ── Uncertainty decomposition visualisation ───────────────────────────────────

fig, axes = plt.subplots(1, 3, figsize=(15, 4.5))
fig.suptitle('Uncertainty decomposition across the test set', fontsize=13)

axes[0].hist(results['predictive_entropy'], bins=20,
             color='#534AB7', alpha=0.7, edgecolor='white')
axes[0].set_xlabel('Total uncertainty (entropy)', fontsize=10)
axes[0].set_ylabel('Count', fontsize=10)
axes[0].set_title('Total uncertainty', fontsize=11)
axes[0].spines[['top', 'right']].set_visible(False)

idx_sorted = results['predictive_entropy'].argsort().values
axes[1].bar(range(len(results)),
            results['epistemic_mi'].values[idx_sorted],
            color='#1D9E75', alpha=0.8, label='Epistemic (MI)')
axes[1].bar(range(len(results)),
            results['aleatoric_h'].values[idx_sorted],
            bottom=results['epistemic_mi'].values[idx_sorted],
            color='#D85A30', alpha=0.8, label='Aleatoric (H)')
axes[1].set_xlabel('Cases sorted by total uncertainty', fontsize=10)
axes[1].set_ylabel('Uncertainty (nats)', fontsize=10)
axes[1].set_title('Epistemic vs aleatoric per case', fontsize=11)
axes[1].legend(fontsize=9)
axes[1].spines[['top', 'right']].set_visible(False)

correct_mask = results['correct'].astype(bool)
axes[2].scatter(results.loc[correct_mask, 'mean_pred'],
                results.loc[correct_mask, 'predictive_entropy'],
                c='#1D9E75', alpha=0.6, s=30, label='Correct')
axes[2].scatter(results.loc[~correct_mask, 'mean_pred'],
                results.loc[~correct_mask, 'predictive_entropy'],
                c='#D85A30', alpha=0.8, s=55, marker='x', label='Wrong')
axes[2].set_xlabel('Mean P(pneumonia)', fontsize=10)
axes[2].set_ylabel('Predictive entropy', fontsize=10)
axes[2].set_title('Confidence vs uncertainty', fontsize=11)
axes[2].legend(fontsize=9)
axes[2].spines[['top', 'right']].set_visible(False)

plt.tight_layout()
plt.show()

# Run stochastic inference on validation set
# We use validation entropy distribution to LEARN the uncertainty threshold.
# Then we apply that threshold to the TEST set.

val_results, _ = stochastic_inference_vi(
    model,
    val_loader,
    T=T_SAMPLES,
    device=device
)

# Fit GMM on VALIDATION entropy

X_val = val_results['predictive_entropy'].values.reshape(-1, 1)

gmm = GaussianMixture(
    n_components=2,
    random_state=42
)

gmm.fit(X_val)

# Find threshold between the two Gaussians

means = gmm.means_.flatten()
means = np.sort(means)

# midpoint between Gaussian means
threshold = means.mean()

print('GMM learned from VALIDATION set')
print(f'Cluster means         : {means}')
print(f'Entropy threshold     : {threshold:.4f}')

# Apply threshold to TEST set

results['uncertainty_group'] = np.where(
    results['predictive_entropy'] > threshold,
    'Uncertain',
    'Certain'
)

certain   = results[results['uncertainty_group'] == 'Certain']
uncertain = results[results['uncertainty_group'] == 'Uncertain']

print()
print(f'Certain group   : {len(certain)} cases')
print(f'Uncertain group : {len(uncertain)} cases')


# Metrics function

def group_metrics(df, name):

    acc = df['correct'].mean()

    auc = (
        roc_auc_score(df['true_label'], df['mean_pred'])
        if df['true_label'].nunique() > 1
        else float('nan')
    )

    mean_ent = df['predictive_entropy'].mean()
    mean_var = df['variance'].mean()

    return {
        'Group'         : name,
        'N Cases'       : len(df),
        'Accuracy'      : acc,
        'AUC-ROC'       : auc,
        'Mean Entropy'  : mean_ent,
        'Mean Variance' : mean_var,
    }


# Compute metrics

certain_metrics   = group_metrics(certain, 'Certain')
uncertain_metrics = group_metrics(uncertain, 'Uncertain')

summary_df = pd.DataFrame([
    certain_metrics,
    uncertain_metrics
])

summary_display = summary_df.copy()

for col in ['Accuracy', 'AUC-ROC', 'Mean Entropy', 'Mean Variance']:
    summary_display[col] = summary_display[col].map(lambda x: f'{x:.4f}')

print('\n📊 Certain vs Uncertain Summary')
display(summary_display)


# Extract values

acc_cert   = certain_metrics['Accuracy']
acc_uncert = uncertain_metrics['Accuracy']

auc_cert   = certain_metrics['AUC-ROC']
auc_uncert = uncertain_metrics['AUC-ROC']


# Mann-Whitney U test (entropy difference between correct vs wrong)

entropy_correct = results.loc[
    results['correct'] == 1,
    'predictive_entropy'
].values

entropy_wrong = results.loc[
    results['correct'] == 0,
    'predictive_entropy'
].values

if len(entropy_wrong) > 0 and len(entropy_correct) > 0:

    stat, p_value = stats.mannwhitneyu(
        entropy_correct,
        entropy_wrong,
        alternative='less'
    )

    print()
    print('📊 Mann-Whitney U test')
    print('H0: entropy(correct) == entropy(wrong)')

    print(f'U statistic : {stat:.2f}')
    print(f'p-value     : {p_value:.4f}')


# Accuracy comparison (z-test)

stat2, p2 = proportions_ztest(
    [
        int(acc_cert * len(certain)),
        int(acc_uncert * len(uncertain))
    ],
    [
        len(certain),
        len(uncertain)
    ]
)

print()
print('📊 Z-test for accuracy difference')
print('H0: accuracy(certain) == accuracy(uncertain)')

print(f'Z statistic : {stat2:.3f}')
print(f'p-value     : {p2:.4f}')

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))

fig.suptitle(
    'Certain vs Uncertain Groups',
    fontsize=13
)

groups = ['Certain', 'Uncertain']
colors = ['#1D9E75', '#D85A30']


# Accuracy comparison

axes[0].bar(groups, [acc_cert, acc_uncert], color=colors, alpha=0.8, width=0.5)

axes[0].set_ylim(0, 1.05)
axes[0].set_ylabel('Accuracy')
axes[0].set_title('Accuracy by group')

for i, v in enumerate([acc_cert, acc_uncert]):
    axes[0].text(i, v + 0.02, f'{v:.3f}', ha='center', fontsize=11, fontweight='bold')

axes[0].spines[['top', 'right']].set_visible(False)


# AUC comparison (descriptive only)

axes[1].bar(groups, [auc_cert, auc_uncert], color=colors, alpha=0.8, width=0.5)

axes[1].set_ylim(0, 1.05)
axes[1].set_ylabel('AUC-ROC')
axes[1].set_title('AUC by group')

for i, v in enumerate([auc_cert, auc_uncert]):
    axes[1].text(i, v + 0.02, f'{v:.3f}', ha='center', fontsize=11, fontweight='bold')

axes[1].spines[['top', 'right']].set_visible(False)


# Entropy distributions

axes[2].hist(
    certain['predictive_entropy'],
    bins=15,
    alpha=0.6,
    color='#1D9E75',
    edgecolor='white',
    label='Certain'
)

axes[2].hist(
    uncertain['predictive_entropy'],
    bins=15,
    alpha=0.6,
    color='#D85A30',
    edgecolor='white',
    label='Uncertain'
)

axes[2].axvline(
    threshold,
    color='black',
    linestyle='--',
    linewidth=1.5,
    label=f'Threshold = {threshold:.3f}'
)

axes[2].set_xlabel('Predictive entropy')
axes[2].set_ylabel('Count')
axes[2].set_title('Entropy distributions')
axes[2].legend()

axes[2].spines[['top', 'right']].set_visible(False)

plt.tight_layout()
plt.show()

# Visualisation helper

LABEL_NAMES = {0: 'Normal', 1: 'Pneumonia'}

def tensor_to_display(x):

    # take first channel
    img = x[0].numpy()

    # rotate 90° clockwise
    img = np.rot90(img, k=-1)

    # mirror flip (left-right flip)
    img = np.fliplr(img)

    return np.clip(img, 0, 1)


def plot_xray_group(indices, results_df, images_list,
                    group_title, n_show=2):

    indices = list(indices)[:n_show]

    if not indices:
        print(f'No cases to show for: {group_title}')
        return

    fig, axes = plt.subplots(
        1,
        len(indices),
        figsize=(5 * len(indices), 5.5)
    )

    if len(indices) == 1:
        axes = [axes]

    fig.suptitle(
        group_title,
        fontsize=13,
        fontweight='bold',
        y=1.02
    )

    for ax, idx in zip(axes, indices):

        row = results_df.loc[idx]

        # now rotated + mirrored
        img = tensor_to_display(images_list[idx])

        ax.imshow(img, cmap='gray')
        ax.axis('off')

        true_name = LABEL_NAMES[int(row['true_label'])]
        pred_name = LABEL_NAMES[int(row['pred_label'])]

        verdict = '✓ Correct' if row['correct'] else '✗ Wrong'

        color = '#1D9E75' if row['correct'] else '#D85A30'

        title = (
            f'{verdict}\n'
            f'True: {true_name}\n'
            f'Pred: {pred_name} ({row["mean_pred"]:.2f})\n'
            f'Entropy: {row["predictive_entropy"]:.4f}\n'
            f'Variance: {row["variance"]:.5f}'
        )

        ax.set_title(
            title,
            fontsize=9,
            color=color,
            pad=5
        )

    plt.tight_layout()
    plt.show()

# MOST CERTAIN CASES

most_certain_idx = (
    certain
    .nsmallest(2, 'predictive_entropy')
    .index.tolist()
)

plot_xray_group(
    most_certain_idx,
    results,
    test_images,
    '🟢 Most Certain Cases'
)

# MOST UNCERTAIN CASES

most_uncertain_idx = (
    uncertain
    .nlargest(2, 'predictive_entropy')
    .index.tolist()
)

plot_xray_group(
    most_uncertain_idx,
    results,
    test_images,
    '🔴 Most Uncertain Cases'
)

results['method'] = 'VI'

save_path = '/kaggle/working/results_vi.csv'
results.to_csv(save_path, index=False)

print(f'Results saved to: {save_path}')