Session 12 — Evidential Deep Learning: 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.nn.functional as F
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)

class EDLDenseNet121(nn.Module):
    """
    MONAI DenseNet121 backbone with an Evidential Deep Learning head.

    The head outputs K evidence values (via ReLU, NOT Softmax).
    Evidence is converted to Dirichlet parameters: alpha_k = e_k + 1.
    Uncertainty is u = K / sum(alpha_k) — high when evidence is low.
    """

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

        backbone = DenseNet121(
            spatial_dims=2,
            in_channels=3,
            out_channels=num_classes,
            pretrained=True,
        )
        self.features = backbone.features
        self.pool     = nn.AdaptiveAvgPool2d((1, 1))

        # Standard linear head — but NO Softmax
        # Output = evidence values (non-negative via ReLU in forward)
        self.evidence_head = nn.Linear(1024, num_classes)

        if freeze_backbone:
            for param in self.features.parameters():
                param.requires_grad = False

    def forward(self, x):
        x = self.features(x)
        x = torch.relu(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)              # [B, 1024]
        evidence = F.relu(self.evidence_head(x))  # [B, K] — non-negative evidence
        return evidence

    def predict_with_uncertainty(self, x):
        """
        Single forward pass → full uncertainty decomposition.
        No T passes needed. All quantities are analytic.
        """
        evidence = self.forward(x)           # [B, K]
        alpha    = evidence + 1              # Dirichlet params [B, K]
        S        = alpha.sum(dim=1, keepdim=True)  # [B, 1]
        K        = self.num_classes

        probs    = alpha / S                 # mean prediction [B, K]
        u        = K / S.squeeze(1)         # total uncertainty [B]

        # Epistemic: sum of Dirichlet variances
        epistemic = (probs * (1 - probs) / (S + 1)).sum(dim=1)   # [B]

        # Aleatoric: residual after removing epistemic
        aleatoric = (probs * (1 - probs)).sum(dim=1) - epistemic  # [B]

        return probs, u, epistemic, aleatoric, alpha


model = EDLDenseNet121(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)
print(f'Total parameters    : {total_params:,}')
print(f'Trainable parameters: {trainable:,}')
print()
print('EDL head: nn.Linear(1024, 2) + ReLU')
print('Same number of parameters as a standard classifier.')
print('No extra parameters vs standard. No T passes at inference.')

def kl_divergence_dirichlet(alpha, K):
    """
    KL divergence between Dir(alpha) and Dir(1,...,1) (uniform prior).
    Computed analytically using lgamma and digamma.
    """
    beta  = torch.ones_like(alpha)         # uniform prior parameters
    S_a   = alpha.sum(dim=1, keepdim=True)
    S_b   = beta.sum(dim=1, keepdim=True)

    kl = (torch.lgamma(S_a) - torch.lgamma(S_b)
          - torch.lgamma(alpha).sum(dim=1, keepdim=True)
          + torch.lgamma(beta).sum(dim=1, keepdim=True)
          + ((alpha - beta) *
             (torch.digamma(alpha) -
              torch.digamma(S_a))).sum(dim=1, keepdim=True))
    return kl.squeeze(1).mean()


def edl_loss(evidence, labels, epoch, num_epochs, K=2, lambda_max=1.0):
    """
    EDL loss = Bayes risk + annealed KL regularisation.

    Args:
        evidence  : [B, K] — raw evidence output (after ReLU)
        labels    : [B]    — integer class labels
        epoch     : current epoch (for annealing)
        num_epochs: total epochs (for annealing)
        K         : number of classes
        lambda_max: maximum KL weight (reached at end of training)
    """
    alpha    = evidence + 1                     # [B, K]
    S        = alpha.sum(dim=1, keepdim=True)   # [B, 1]
    y_onehot = F.one_hot(labels, num_classes=K).float()  # [B, K]

    # ── Bayes risk (fit term) ─────────────────────────────────────────────────
    # Expected MSE under the Dirichlet
    p_hat = alpha / S                           # [B, K]
    err   = (y_onehot - p_hat) ** 2            # [B, K]
    var   = p_hat * (1 - p_hat) / (S + 1)     # [B, K]
    fit   = (err + var).sum(dim=1).mean()

    # ── KL regularisation (annealed) ──────────────────────────────────────────
    # Remove evidence on correct class before computing KL
    # so we only penalise evidence on wrong classes
    alpha_tilde = y_onehot + (1 - y_onehot) * alpha
    lambda_t    = lambda_max * min(1.0, epoch / (num_epochs * 0.5))
    kl          = kl_divergence_dirichlet(alpha_tilde, K)

    return fit + lambda_t * kl, fit.item(), kl.item(), lambda_t


print('EDL loss implemented:')
print('  Bayes risk (fit term) + λ_t × KL(Dir(α̃) || Dir(1))')
print('  λ_t anneals from 0 to lambda_max over the first 50% of training')

# Training setup

NUM_EPOCHS = 50
LAMBDA_MAX = 0.5

FINAL_PATH  = "/kaggle/working/final_model.pth"
LATEST_PATH = "/kaggle/working/latest_model.pth"

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

print(f'NUM_EPOCHS  : {NUM_EPOCHS}')
print(f'FINAL_PATH  : {FINAL_PATH}')
print(f'LATEST_PATH : {LATEST_PATH}')

# History

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


# Training loop

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

    # ─────────────── Train ───────────────
    model.train()

    train_loss = 0.0
    train_fit  = 0.0
    train_kl   = 0.0
    lam_t      = 0.0

    for batch in tqdm(train_loader, desc=f'Epoch {epoch}/{NUM_EPOCHS}', leave=False):

        images = batch['image'].to(device, dtype=torch.float32)
        labels = batch['label'].to(device, dtype=torch.long)

        optimizer.zero_grad()

        evidence = model(images)

        loss, fit, kl, lam_t = edl_loss(
            evidence,
            labels,
            epoch,
            NUM_EPOCHS,
            K=2,
            lambda_max=LAMBDA_MAX
        )

        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        train_fit  += fit
        train_kl   += kl


    # ─────────────── 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 = batch['label'].to(device, dtype=torch.long)

            evidence = model(images)

            loss, _, _, _ = edl_loss(
                evidence,
                labels,
                epoch,
                NUM_EPOCHS,
                K=2,
                lambda_max=LAMBDA_MAX
            )

            val_loss += loss.item()

            probs, _, _, _, _ = model.predict_with_uncertainty(images)

            val_probs.extend(probs[:, 1].cpu().tolist())
            val_preds.extend(probs.argmax(1).cpu().tolist())
            val_labels.extend(labels.cpu().tolist())


    # ─────────────── Metrics ───────────────

    train_loss /= len(train_loader)
    train_fit  /= len(train_loader)
    train_kl   /= len(train_loader)
    val_loss   /= len(val_loader)

    val_acc = accuracy_score(val_labels, val_preds)
    val_auc = roc_auc_score(val_labels, val_probs)


    # ─────────────── History ───────────────

    history['train_loss'].append(train_loss)
    history['fit_loss'].append(train_fit)
    history['kl_loss'].append(train_kl)
    history['val_loss'].append(val_loss)
    history['lambda_t'].append(lam_t)
    history['val_acc'].append(val_acc)
    history['val_auc'].append(val_auc)


    # ─────────────── Logging ───────────────

    print(
    f'Epoch {epoch:02d}/{NUM_EPOCHS} | '
    f'Train Loss: {train_loss:.4f} | '
    f'Val Loss: {val_loss:.4f} | '
    f'Fit: {train_fit:.4f} | '
    f'KL: {train_kl:.4f} | '
    f'λ: {lam_t:.2f} | '
    f'Val Acc: {val_acc:.3f} | '
    f'Val AUC: {val_auc:.3f}' 
    )


    # ─────────────── Save latest model ───────────────

    torch.save(model.state_dict(), LATEST_PATH)
    print(f"📌 Latest model saved at: {LATEST_PATH}")


# Save final model

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

print("\n🏁 Training complete")
print(f"📌 Final model saved at:  {FINAL_PATH}")

# Training curves — note we track fit and KL separately
fig, axes = plt.subplots(1, 5, figsize=(15, 8))
fig.suptitle('EDL DenseNet121 — Training curves', fontsize=13)
epochs = range(1, len(history['train_loss']) + 1)

for ax, key, color, title in zip(
    axes.flatten(),
    ['train_loss', 'fit_loss', 'kl_loss', 'lambda_t', 'val_loss'],
    ['#534AB7', '#D85A30', '#1D9E75', '#E8A838', '#1D9E75', '#7F77DD'],
    ['Total EDL Loss', 'Fit term (Bayes risk)', 'KL term',
     'KL weight λ_t (annealing)', 'Validation Loss']
):
    ax.plot(epochs, history[key], color=color, linewidth=2)
    ax.set_title(title, fontsize=11)
    ax.set_xlabel('Epoch')
    ax.spines[['top', 'right']].set_visible(False)

plt.tight_layout()
plt.show()

# Loading the best model

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

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

BEST_PATH = "/kaggle/input/models/benyamingheiji/edl-final-model/pytorch/default/1/edl_final_model.pth"

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

model.eval()

print('✅ Best model loaded successfully')
print(f'Loaded weights from: {BEST_PATH}')

def edl_inference(model, loader, device=device):
    """
    EDL inference: ONE deterministic forward pass per image.

    Uncertainty decomposition (Option 2 — Dirichlet Statistics):
        predictive_entropy = entropy of mean prediction         (total)
        aleatoric          = expected entropy under Dirichlet   (data uncertainty)
        epistemic          = predictive_entropy - aleatoric     (model uncertainty, MI)
        vacuity            = K/S                                (EDL-specific, separate)

    predictive_entropy = epistemic + aleatoric  ✅
    """
    model.eval()

    all_mean_preds, all_vacuity, all_epistemic, all_aleatoric, all_predictive = [], [], [], [], []
    all_evidence_sum, all_labels, all_paths, all_images = [], [], [], []

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

            probs, u, epistemic, aleatoric, alpha = \
                model.predict_with_uncertainty(images)

            S                  = alpha.sum(dim=1)
            predictive_entropy = epistemic + aleatoric   # [B] additive ✅

            all_mean_preds.extend(probs[:, 1].cpu().tolist())
            all_vacuity.extend(u.cpu().tolist())
            all_predictive.extend(predictive_entropy.cpu().tolist())
            all_epistemic.extend(epistemic.cpu().tolist())
            all_aleatoric.extend(aleatoric.cpu().tolist())
            all_evidence_sum.extend(S.cpu().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(),
        'vacuity'            : all_vacuity,      # K/S — EDL-specific
        'predictive_entropy' : all_predictive,   # epistemic + aleatoric ✅
        'epistemic_mi'       : all_epistemic,
        'aleatoric_h'        : all_aleatoric,
        'evidence_strength'  : all_evidence_sum,
        'correct'            : correct.tolist(),
    })
    return results, all_images


# ── Inference ──────────────────────────────────────────────────────────────────

results, test_images = edl_inference(model, 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('\nUncertainty decomposition (mean across test set):')
print(f'  Predictive entropy : {results["predictive_entropy"].mean():.6f}')
print(f'  Epistemic (MI)     : {results["epistemic_mi"].mean():.6f}')
print(f'  Aleatoric (H)      : {results["aleatoric_h"].mean():.6f}')
print(f'  Vacuity (K/S)      : {results["vacuity"].mean():.4f}')
print(f'  Mean evidence S    : {results["evidence_strength"].mean():.4f}')

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

axes[0].hist(results['predictive_entropy'], bins=20,
             color='#534AB7', alpha=0.7, edgecolor='white')
axes[0].set_xlabel('Predictive entropy (epistemic + aleatoric)', fontsize=10)
axes[0].set_ylabel('Count', fontsize=10)
axes[0].set_title('Total predictive entropy', 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 predictive entropy', fontsize=10)
axes[1].set_ylabel('Uncertainty', 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, _ = edl_inference(
    model,
    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(f'\nCertain 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_epi = df['epistemic_mi'].mean()
    mean_ale = df['aleatoric_h'].mean()
    mean_vac = df['vacuity'].mean()

    return {
        'Group'             : name,
        'N Cases'           : len(df),
        'Accuracy'          : acc,
        'AUC-ROC'           : auc,
        'Mean Pred Entropy' : mean_ent,
        'Mean Epistemic'    : mean_epi,
        'Mean Aleatoric'    : mean_ale,
        'Mean Vacuity'      : mean_vac,
    }


# ── Summary table ─────────────────────────────────────────────────────────────

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 Pred Entropy',
            'Mean Epistemic', 'Mean Aleatoric', 'Mean Vacuity']:
    summary_display[col] = summary_display[col].map(lambda x: f'{x:.4f}')

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

acc_cert   = certain_metrics['Accuracy']
acc_uncert = uncertain_metrics['Accuracy']
auc_cert   = certain_metrics['AUC-ROC']
auc_uncert = uncertain_metrics['AUC-ROC']


# ── Statistical tests ─────────────────────────────────────────────────────────

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('\n📊 Mann-Whitney U test')
    print('H0: entropy(correct) == entropy(wrong)')
    print(f'U statistic : {stat:.2f}')
    print(f'p-value     : {p_value:.4f}')

stat2, p2 = proportions_ztest(
    [int(acc_cert * len(certain)), int(acc_uncert * len(uncertain))],
    [len(certain), len(uncertain)]
)
print('\n📊 Z-test for accuracy difference')
print('H0: accuracy(certain) == accuracy(uncertain)')
print(f'Z statistic : {stat2:.3f}')
print(f'p-value     : {p2:.4f}')


# ── Certain vs Uncertain plots ────────────────────────────────────────────────

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
fig.suptitle('Certain vs Uncertain Groups', fontsize=13)

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

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)

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)

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):
    img = x[0].numpy()
    img = np.rot90(img, k=-1)
    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]
        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'Epistemic: {row["epistemic_mi"]:.5f}\n'
            f'Aleatoric: {row["aleatoric_h"]:.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'] = 'EDL'

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

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