Session 10 — Deep Ensembles: 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
from torch_uncertainty.metrics.classification import CalibrationError, AURC

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

M          = 3       # number of ensemble members
NUM_EPOCHS = 20

def build_member(seed, num_classes=2, freeze_backbone=True):
    """
    Build one ensemble member — a standard DenseNet121 with a
    fresh classification head seeded differently per member.

    No BayesLinear. No dropout trick. Just a standard model.
    The diversity comes from different random seeds + training stochasticity.
    """
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    backbone = DenseNet121(
        spatial_dims=2,
        in_channels=3,
        out_channels=num_classes,
        pretrained=True,
    )

    # Reinitialise the classification head with this member's seed
    # (backbone stays at ImageNet pretrained weights)
    nn.init.kaiming_normal_(backbone.class_layers.out.weight)
    nn.init.zeros_(backbone.class_layers.out.bias)

    if freeze_backbone:
        for name, param in backbone.named_parameters():
            if 'class_layers' not in name:
                param.requires_grad = False

    return backbone.to(device)


# Build M members with different seeds
member_seeds = [42, 123, 777]
ensemble     = [build_member(seed) for seed in member_seeds]

print(f'Ensemble size: M = {M} members')
trainable = sum(p.numel() for p in ensemble[0].parameters() if p.requires_grad)
print(f'Trainable params per member: {trainable:,}')
print(f'Total trainable params (M×): {trainable * M:,}')
print()
print('Compare with Session 5 (VI): 2 × head params (mu + sigma)')
print('Compare with Session 6 (MC Dropout): 1 × head params (standard)')
print('Deep Ensembles: M × head params — cost scales with M')

# Deep Ensemble Training

M          = 3
NUM_EPOCHS = 100

EARLY_STOP_PATIENCE = 10
PLATEAU_PATIENCE    = 3
PLATEAU_FACTOR      = 0.5
MIN_LR              = 1e-6

loss_fn = nn.CrossEntropyLoss()

all_history      = []
best_model_paths = []

# Train each ensemble member independently

for m_idx, (model, seed) in enumerate(zip(ensemble, member_seeds)):

    print(f'\n{"="*60}')
    print(f'Training member {m_idx+1}/{M}  (seed={seed})')
    print(f'{"="*60}')

    # ---------------------------------------------------------------------
    # Reproducibility
    # ---------------------------------------------------------------------
    random.seed(seed)

    np.random.seed(seed)

    torch.manual_seed(seed)

    torch.cuda.manual_seed_all(seed)

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

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

    # ---------------------------------------------------------------------
    # Best-model tracking
    # ---------------------------------------------------------------------
    best_val_loss = np.inf

    patience_counter = 0

    best_path = (
        f'/kaggle/working/best_member_{m_idx+1}.pth'
    )

    best_model_paths.append(best_path)

    # ---------------------------------------------------------------------
    # History
    # ---------------------------------------------------------------------
    history = {
        'train_loss': [],
        'val_loss': [],
        'val_acc': [],
        'val_auc': [],
        'lr': [],
    }

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

            logits = model(images)

            loss = loss_fn(logits, labels)

            loss.backward()

            optimizer.step()

            epoch_loss += loss.item()

        avg_train_loss = epoch_loss / len(train_loader)

        # -----------------------------------------------------------------
        # VALIDATION
        # -----------------------------------------------------------------
        model.eval()

        val_loss = 0.0

        val_probs  = []
        val_preds  = []
        val_labels = []

        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
                )

                logits = model(images)

                loss = loss_fn(logits, labels_tensor)

                val_loss += loss.item()

                probs = torch.softmax(
                    logits,
                    dim=-1
                ).cpu().numpy()

                labels = labels_tensor.cpu().numpy()

                val_probs.extend(probs[:, 1].tolist())

                val_preds.extend(
                    probs.argmax(axis=1).tolist()
                )

                val_labels.extend(labels.tolist())

        avg_val_loss = val_loss / len(val_loader)

        val_acc = accuracy_score(
            val_labels,
            val_preds
        )

        val_auc = roc_auc_score(
            val_labels,
            val_probs
        )

        # -----------------------------------------------------------------
        # Scheduler update
        # -----------------------------------------------------------------
        scheduler.step(avg_val_loss)

        current_lr = optimizer.param_groups[0]['lr']

        # -----------------------------------------------------------------
        # Save metrics
        # -----------------------------------------------------------------
        history['train_loss'].append(avg_train_loss)

        history['val_loss'].append(avg_val_loss)

        history['val_acc'].append(val_acc)

        history['val_auc'].append(val_auc)

        history['lr'].append(current_lr)

        # -----------------------------------------------------------------
        # Save best model
        # -----------------------------------------------------------------
        if avg_val_loss < best_val_loss:

            best_val_loss = avg_val_loss

            patience_counter = 0

            torch.save(
                model.state_dict(),
                best_path
            )

            print(
                f'✅ Saved best model '
                f'(val_loss={best_val_loss:.4f})'
            )

        else:

            patience_counter += 1

        # -----------------------------------------------------------------
        # Logging
        # -----------------------------------------------------------------
        print(
            f'Epoch {epoch:02d}/{NUM_EPOCHS} | '
            f'Train Loss: {avg_train_loss:.4f} | '
            f'Val Loss: {avg_val_loss:.4f} | '
            f'Val Acc: {val_acc:.3f} | '
            f'Val AUC: {val_auc:.3f} | '
            f'LR: {current_lr:.6f}'
        )

        # -----------------------------------------------------------------
        # Early stopping
        # -----------------------------------------------------------------
        if patience_counter >= EARLY_STOP_PATIENCE:

            print()

            print(
                f'🛑 Early stopping triggered '
                f'for member {m_idx+1}'
            )

            break

    # ---------------------------------------------------------------------
    # Save history
    # ---------------------------------------------------------------------
    all_history.append(history)

print('\n🎉 All ensemble members trained!')

# ── Train & validation loss curves (separate figure per member) ─────────────

colors_m = ['#7F77DD', '#1D9E75', '#D85A30']

for m_idx, (h, color) in enumerate(zip(all_history, colors_m)):

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

    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))

    # ── Train loss (left) ─────────────────────────────
    axes[0].plot(
        epochs,
        h['train_loss'],
        color=color,
        linewidth=2
    )
    axes[0].set_title(f'Member {m_idx+1} - Train Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].spines[['top', 'right']].set_visible(False)

    # ── Validation loss (right) ───────────────────────
    axes[1].plot(
        epochs,
        h['val_loss'],
        color=color,
        linewidth=2
    )
    axes[1].set_title(f'Member {m_idx+1} - Val Loss')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Loss')
    axes[1].spines[['top', 'right']].set_visible(False)

    plt.tight_layout()
    plt.show()

# Load BEST weights for all ensemble members

# I renamed the best models and added them to inputs

model_paths = [
    '/kaggle/input/models/benyamingheiji/de-best-model1/pytorch/default/1/DE_best_model1.pth',
    '/kaggle/input/models/benyamingheiji/de-best-model2/pytorch/default/1/DE_best_model2.pth',
    '/kaggle/input/models/benyamingheiji/de-best-model3/pytorch/default/1/DE_best_model3.pth'
]

for m_idx, (model, path) in enumerate(zip(ensemble, model_paths)):

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

    model.eval()

    print(
        f'✅ Loaded best weights for '
        f'Member {m_idx+1} from {path}'
    )

print('\n🎯 All best ensemble members restored!')

def ensemble_inference(ensemble, loader, device=device):
    """
    Ensemble inference: one deterministic forward pass per member per image.

    Key difference from VI and MC Dropout:
    - No stochastic sampling during inference
    - Each member is fully deterministic (model.eval(), no dropout)
    - Uncertainty = disagreement between M independent predictions

    Returns results DataFrame and raw image tensors.
    """
    for m in ensemble:
        m.eval()

    all_mean_preds, all_entropy, all_variance = [], [], []
    all_epistemic, all_aleatoric              = [], []
    all_labels, all_paths, all_images         = [], [], []
    # Per-member predictions for diversity analysis
    all_member_preds = [[] for _ in ensemble]

    with torch.no_grad():
        for batch in tqdm(loader, desc='Ensemble 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))

            # One deterministic pass per member
            member_probs = []
            for m_idx, m in enumerate(ensemble):
                probs = torch.softmax(m(images), dim=-1).cpu().numpy()  # [B, 2]
                member_probs.append(probs)
                all_member_preds[m_idx].extend(probs[:, 1].tolist())

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

            # Total uncertainty: predictive entropy of the mean
            total_H = -np.sum(
                mean_probs * np.log(mean_probs + 1e-8), axis=-1
            )

            # Aleatoric: average entropy per member
            per_H       = -np.sum(
                member_probs * np.log(member_probs + 1e-8), axis=-1
            )  # [M, B]
            aleatoric_H = per_H.mean(axis=0)

            # Epistemic: mutual information
            epistemic_MI = total_H - aleatoric_H

            # Variance of P(pneumonia) across members
            variance = member_probs[:, :, 1].var(axis=0)

            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, all_member_preds


results, test_images, member_preds = ensemble_inference(
    ensemble, test_loader, 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}')

print()
print('Per-member performance:')

for m_idx, preds_m in enumerate(member_preds):

    preds_m   = np.array(preds_m)

    preds_bin = (preds_m >= 0.5).astype(int)

    acc_m = (
        preds_bin == results['true_label'].values
    ).mean()

    auc_m = roc_auc_score(
        results['true_label'],
        preds_m
    )

    print(
        f'  Member {m_idx+1} '
        f'| Accuracy: {acc_m:.3f} '
        f'| AUC: {auc_m:.3f}'
    )

print()
print(
    f'Ensemble '
    f'| Accuracy: {results["correct"].mean():.3f} '
    f'| AUC: {roc_auc_score(results["true_label"], results["mean_pred"]):.3f}'
)

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 (Deep Ensembles)', fontsize=13)

# Total entropy
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)

# Epistemic vs aleatoric per case
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 vs wrong
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, _, _ = ensemble_inference(
    ensemble,
    val_loader,
    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'] = 'Deep_Ensembles'

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

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