VGG

# Standard libraries
import time
from typing import Any, List

# 3pps
# Third-party libraries
import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.manifold import TSNE
from torch import nn
from torch.utils.data import DataLoader
from torchinfo import summary
from torchvision import datasets, transforms
from tqdm import tqdm

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA device: {torch.cuda.get_device_name(0)}")

# Global configuration
BATCH_SIZE: int = 128
NUM_EPOCHS: int = 1
LEARNING_RATE: float = 1e-3
WEIGHT_DECAY: float = 5e-4
NUM_CLASSES: int = 10  # CIFAR-10 has 10 classes
INPUT_SIZE: int = 32  # CIFAR-10: 32×32 images

# CIFAR-10 class names
CIFAR10_CLASSES = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
]

print("Configuration:")
print(f"  Batch size: {BATCH_SIZE}")
print(f"  Epochs: {NUM_EPOCHS}")
print(f"  Learning rate: {LEARNING_RATE}")
print(f"  Number of classes: {NUM_CLASSES}")

def show_images(images, labels, predictions=None, classes=CIFAR10_CLASSES):
    """
    Visualizes a set of images with their labels.

    Args:
        images: Image tensor [N, C, H, W]
        labels: Label tensor [N]
        predictions: Optional prediction tensor [N]
        classes: List of class names
    """
    n_images = len(images)
    fig, axes = plt.subplots(1, n_images, figsize=(2 * n_images, 3))

    if n_images == 1:
        axes = [axes]

    for idx, (img, label, ax) in enumerate(zip(images, labels, axes)):
        # Denormalize image for visualization
        img = img / 2 + 0.5  # Revert normalization [-1, 1] -> [0, 1]
        img = img.numpy().transpose((1, 2, 0))

        ax.imshow(img)

        title = f"True: {classes[label]}"
        if predictions is not None:
            pred = predictions[idx]
            color = "green" if pred == label else "red"
            title += f"\nPred: {classes[pred]}"
            ax.set_title(title, fontsize=10, color=color, fontweight="bold")
        else:
            ax.set_title(title, fontsize=10, fontweight="bold")

        ax.axis("off")

    plt.tight_layout()
    plt.show()

print("Visualization function defined correctly")

from torch.utils.data import Subset

# Transformations with data augmentation for training
transform_train = transforms.Compose(
    [
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)
        ),
    ]
)

# Transformations for validation (no augmentation)
transform_test = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ]
)

print("Downloading CIFAR-10 dataset...")

# Training set
train_dataset_full = datasets.CIFAR10(
    root="./data", train=True, download=True, transform=transform_train
)

# Test set
test_dataset_full = datasets.CIFAR10(
    root="./data", train=False, download=True, transform=transform_test
)

# Limit samples
train_dataset = Subset(train_dataset_full, range(5000))
test_dataset = Subset(test_dataset_full, range(1000))

print("\nDataset statistics:")
print(f"  Training samples: {len(train_dataset):,}")
print(f"  Test samples: {len(test_dataset):,}")
print(f"  Number of classes: {len(train_dataset_full.classes)}")
print(f"  Classes: {', '.join(train_dataset_full.classes)}")
print("  Image size: 32×32 pixels (RGB)")

print("Configuring DataLoaders:")
print(f"  Batch size: {BATCH_SIZE}")

# Training DataLoader
train_dataloader = DataLoader(
    dataset=train_dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True,
)

# Test DataLoader
test_dataloader = DataLoader(
    dataset=test_dataset,
    batch_size=BATCH_SIZE,
    shuffle=False,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True,
)

print(f"  Training batches: {len(train_dataloader)}")
print(f"  Test batches: {len(test_dataloader)}")

# Obtain one data batch
data_iter = iter(train_dataloader)
train_images, train_labels = next(data_iter)

print("Batch dimensions:")
print(f"  Images: {train_images.shape}")
print(f"  Labels: {train_labels.shape}")
print(f"\n  Interpretation: {BATCH_SIZE} RGB images of size 32×32 pixels")

# Visualize first 8 examples
print("\nVisualizing first 8 samples...")
show_images(train_images[:8], train_labels[:8])

# Statistics of normalized images
print("\nStatistics after normalization:")
print(f"  Min value: {train_images.min():.3f}")
print(f"  Max value: {train_images.max():.3f}")
print("  Mean per channel:")
for i, channel in enumerate(["R", "G", "B"]):
    print(f"    {channel}: {train_images[:, i, :, :].mean():.3f}")

class VGG16(nn.Module):
    """
    Implementation of VGG-16 adapted for CIFAR-10 (32×32 pixels).

    Architecture:
    - 5 convolutional blocks with configuration [64, 128, 256, 512, 512]
    - All filters are 3×3
    - 2×2 MaxPooling after each block
    - 3 fully connected layers at the end
    - BatchNorm to stabilize training
    - Dropout for regularization
    """

    def __init__(self, num_classes: int = 10, **kwargs: Any) -> None:
        super().__init__(**kwargs)

        self.num_classes = num_classes

        # Block 1: 2 conv layers with 64 filters
        self.block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 32×32 -> 16×16
        )

        # Block 2: 2 conv layers with 128 filters
        self.block2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 16×16 -> 8×8
        )

        # Block 3: 3 conv layers with 256 filters
        self.block3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 8×8 -> 4×4
        )

        # Block 4: 3 conv layers with 512 filters
        self.block4 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 4×4 -> 2×2
        )

        # Block 5: 3 conv layers with 512 filters
        self.block5 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 2×2 -> 1×1
        )

        # Classifier layers
        # For CIFAR-10 (32×32), after 5 poolings: 32 / 2^5 = 1
        # Therefore: 512 × 1 × 1 = 512 features
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 1 * 1, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes),
        )

        # Weight initialization
        self._initialize_weights()

    def _initialize_weights(self):
        """
        Initializes weights using He initialization for layers with ReLU.
        """
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of VGG-16.

        Args:
            x: Input tensor [B, 3, 32, 32]

        Returns:
            Classification logits [B, num_classes]
        """
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = self.block5(x)
        x = self.classifier(x)
        return x

    def get_features(self, x: torch.Tensor) -> torch.Tensor:
        """
        Extracts features before the final classification layer.
        Useful for embedding visualization.
        """
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = self.block5(x)
        x = torch.flatten(x, 1)
        return x

print("VGG-16 architecture defined correctly")

# Create the model
model = VGG16(num_classes=NUM_CLASSES)

# Determine available device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

print(f"Device used: {device}")
print(f"\n{'=' * 70}")
print("VGG-16 ARCHITECTURE SUMMARY")
print(f"{'=' * 70}\n")

# Detailed architecture summary
summary(model, input_size=(BATCH_SIZE, 3, 32, 32), device=str(device))

# Parameter count per block
print(f"\n{'=' * 70}")
print("PARAMETER ANALYSIS PER BLOCK")
print(f"{'=' * 70}")

def count_parameters(module):
    return sum(p.numel() for p in module.parameters())

print(f"  Block 1: {count_parameters(model.block1):>12,} parameters")
print(f"  Block 2: {count_parameters(model.block2):>12,} parameters")
print(f"  Block 3: {count_parameters(model.block3):>12,} parameters")
print(f"  Block 4: {count_parameters(model.block4):>12,} parameters")
print(f"  Block 5: {count_parameters(model.block5):>12,} parameters")
print(f"  Classifier: {count_parameters(model.classifier):>12,} parameters")
print(f"  {'-' * 66}")

total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"  TOTAL: {total_params:>12,} parameters")
print(f"  Trainable: {trainable_params:>12,} parameters")
print(f"  Memory (float32): {total_params * 4 / (1024 ** 2):>10.2f} MB")

print("TRAINING CONFIGURATION")
print(f"{'=' * 70}")
print(f"  Epochs: {NUM_EPOCHS}")
print(f"  Initial learning rate: {LEARNING_RATE}")
print(f"  Weight decay (L2): {WEIGHT_DECAY}")
print(f"  Batch size: {BATCH_SIZE}")
print(f"{'=' * 70}\n")

# Optimizer: SGD with momentum
optimizer = torch.optim.SGD(
    params=model.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=WEIGHT_DECAY
)

# Learning rate scheduler: reduces LR when progress plateaus
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="max",  # Monitor accuracy (to be maximized)
    factor=0.1,  # Reduce LR to 10 % of current value
    patience=3,  # After 3 epochs without improvement
    min_lr=1e-6,
)

# Loss function: Cross-Entropy
loss_function = nn.CrossEntropyLoss()

print("Optimizer: SGD with momentum=0.9")
print("  SGD with momentum accumulates gradients with exponential decay")
print("  This helps escape local minima and accelerates convergence")
print("\nScheduler: ReduceLROnPlateau")
print("  Reduces the learning rate when accuracy stops improving")
print("  Reduction factor: 0.1 (LR → 0.1 × LR)")
print("  Patience: 3 epochs")
print("\nLoss function: CrossEntropyLoss")

from tqdm import tqdm
import time

# Lists to store metrics
train_losses, train_accuracies = [], []
test_losses, test_accuracies = [], []
learning_rates = []

# Auxiliary function to compute accuracy
def calculate_accuracy(outputs, labels):
    _, predicted = torch.max(outputs, 1)
    correct = (predicted == labels).sum().item()
    total = labels.size(0)
    return correct, total

print("STARTING TRAINING\n")
print(f"{'=' * 70}\n")

# Start time
start_time = time.time()

for epoch in range(NUM_EPOCHS):
    epoch_start_time = time.time()

    # ============ TRAINING PHASE ============
    model.train()
    running_loss, correct, total = 0.0, 0, 0

    train_loop = tqdm(
        train_dataloader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS} [TRAIN]", leave=False
    )

    for batch_image, batch_label in train_loop:
        batch_image = batch_image.to(device)
        batch_label = batch_label.to(device)

        optimizer.zero_grad()

        outputs = model(batch_image)
        loss = loss_function(outputs, batch_label)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        batch_correct, batch_total = calculate_accuracy(outputs, batch_label)
        correct += batch_correct
        total += batch_total

        train_loop.set_postfix(
            {"loss": f"{loss.item():.4f}", "acc": f"{100 * correct / total:.2f}%"}
        )

    epoch_train_loss = running_loss / len(train_dataloader)
    epoch_train_acc = 100 * correct / total
    train_losses.append(epoch_train_loss)
    train_accuracies.append(epoch_train_acc)

    # ============ VALIDATION PHASE ============
    model.eval()
    test_loss, correct_test, total_test = 0.0, 0, 0

    test_loop = tqdm(
        test_dataloader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS} [TEST]", leave=False
    )

    with torch.no_grad():
        for images, labels in test_loop:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = loss_function(outputs, labels)

            test_loss += loss.item()
            batch_correct, batch_total = calculate_accuracy(outputs, labels)
            correct_test += batch_correct
            total_test += batch_total

            test_loop.set_postfix(
                {
                    "loss": f"{loss.item():.4f}",
                    "acc": f"{100 * correct_test / total_test:.2f}%",
                }
            )

    epoch_test_loss = test_loss / len(test_dataloader)
    epoch_test_acc = 100 * correct_test / total_test
    test_losses.append(epoch_test_loss)
    test_accuracies.append(epoch_test_acc)

    # Update learning rate scheduler
    scheduler.step(epoch_test_acc)
    current_lr = optimizer.param_groups[0]["lr"]
    learning_rates.append(current_lr)

    # Epoch time
    epoch_time = time.time() - epoch_start_time

    # Epoch report
    print(f"Epoch [{epoch + 1}/{NUM_EPOCHS}] - Time: {epoch_time:.2f}s")
    print(f"  Train → Loss: {epoch_train_loss:.4f} | Acc: {epoch_train_acc:.2f}%")
    print(f"  Test  → Loss: {epoch_test_loss:.4f} | Acc: {epoch_test_acc:.2f}%")
    print(f"  LR: {current_lr:.6f}")
    print(f"  {'─' * 66}\n")

# Total time
total_time = time.time() - start_time

print(f"\n{'=' * 70}")
print("TRAINING COMPLETED")
print(f"{'=' * 70}")
print(f"  Total time: {total_time / 60:.2f} minutes")
print(f"  Average time per epoch: {total_time / NUM_EPOCHS:.2f} seconds")
print(f"  Final accuracy (train): {train_accuracies[-1]:.2f}%")
print(f"  Final accuracy (test): {test_accuracies[-1]:.2f}%")
print(f"  Best accuracy (test): {max(test_accuracies):.2f}%")

# Save the model
torch.save(
    {
        "epoch": NUM_EPOCHS,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "train_losses": train_losses,
        "train_accuracies": train_accuracies,
        "test_losses": test_losses,
        "test_accuracies": test_accuracies,
    },
    "vgg16_cifar10.pth",
)

print("\nModel saved as 'vgg16_cifar10.pth'")

import os

os.makedirs("results", exist_ok=True)

epochs_range = range(1, NUM_EPOCHS + 1)

# Create figure with three subplots
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 5))

# Loss plot
ax1.plot(
    epochs_range, train_losses, "o-", label="Train Loss", linewidth=2, markersize=6
)
ax1.plot(epochs_range, test_losses, "s-", label="Test Loss", linewidth=2, markersize=6)
ax1.set_xlabel("Epoch", fontsize=12, fontweight="bold")
ax1.set_ylabel("Loss", fontsize=12, fontweight="bold")
ax1.set_title("Loss Evolution", fontsize=14, fontweight="bold")
ax1.legend(fontsize=11)
ax1.grid(True, alpha=0.3)
ax1.set_xticks(list(epochs_range))

# Accuracy plot
ax2.plot(
    epochs_range,
    train_accuracies,
    "o-",
    label="Train Accuracy",
    linewidth=2,
    markersize=6,
)
ax2.plot(
    epochs_range,
    test_accuracies,
    "s-",
    label="Test Accuracy",
    linewidth=2,
    markersize=6,
)
ax2.set_xlabel("Epoch", fontsize=12, fontweight="bold")
ax2.set_ylabel("Accuracy (%)", fontsize=12, fontweight="bold")
ax2.set_title("Accuracy Evolution", fontsize=14, fontweight="bold")
ax2.legend(fontsize=11)
ax2.grid(True, alpha=0.3)
ax2.set_xticks(list(epochs_range))

# Learning rate plot
ax3.plot(epochs_range, learning_rates, "o-", color="red", linewidth=2, markersize=6)
ax3.set_xlabel("Epoch", fontsize=12, fontweight="bold")
ax3.set_ylabel("Learning Rate", fontsize=12, fontweight="bold")
ax3.set_title("Learning Rate Schedule", fontsize=14, fontweight="bold")
ax3.set_yscale("log")
ax3.grid(True, alpha=0.3)
ax3.set_xticks(list(epochs_range))

plt.tight_layout()
plt.savefig("results/vgg16_training_history.png", dpi=300, bbox_inches="tight")
plt.show()

# Quantitative analysis
print("\nResult analysis:")
diff = train_accuracies[-1] - test_accuracies[-1]
print(f"  Overfitting detected: {'YES' if diff > 10 else 'NO'}")
print(f"  Train-test gap: {diff:.2f}%")
print(f"  Best epoch (test acc): {np.argmax(test_accuracies) + 1}")
print(f"  Gain since epoch 1: {test_accuracies[-1] - test_accuracies[0]:.2f}%")

# 3pps
import os
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix

os.makedirs("results", exist_ok=True)

print("Generating confusion matrix...")

# Obtain all predictions
model.eval()
all_predictions = []
all_labels = []

with torch.no_grad():
    for images, labels in tqdm(test_dataloader, desc="Evaluating"):
        images = images.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

        all_predictions.extend(predicted.cpu().numpy())
        all_labels.extend(labels.numpy())

# Compute confusion matrix
cm = confusion_matrix(all_labels, all_predictions)

# Visualize confusion matrix
plt.figure(figsize=(12, 10))
sns.heatmap(
    cm,
    annot=True,
    fmt="d",
    cmap="Blues",
    xticklabels=CIFAR10_CLASSES,
    yticklabels=CIFAR10_CLASSES,
    cbar_kws={"label": "Number of samples"},
)
plt.xlabel("Prediction", fontsize=12, fontweight="bold")
plt.ylabel("True Label", fontsize=12, fontweight="bold")
plt.title("Confusion Matrix - VGG16 on CIFAR-10", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.savefig("results/vgg16_confusion_matrix.png", dpi=300, bbox_inches="tight")
plt.show()

# Classification report
print("\nClassification Report:")
print("=" * 70)
print(
    classification_report(
        all_labels, all_predictions, target_names=CIFAR10_CLASSES, digits=3
    )
)

# Per-class accuracy analysis
print("\nPer-Class Accuracy Analysis:")
print("=" * 70)
class_correct = cm.diagonal()
class_total = cm.sum(axis=1)
class_accuracy = 100 * class_correct / class_total

for idx, class_name in enumerate(CIFAR10_CLASSES):
    print(
        f"  {class_name:12s}: {class_accuracy[idx]:6.2f}% "
        f"({class_correct[idx]:4d}/{class_total[idx]:4d})"
    )

# Most confused class pairs
print("\nMost Confused Class Pairs:")
print("=" * 70)
confusion_pairs = []
for i in range(len(CIFAR10_CLASSES)):
    for j in range(len(CIFAR10_CLASSES)):
        if i != j:
            confusion_pairs.append((cm[i, j], CIFAR10_CLASSES[i], CIFAR10_CLASSES[j]))

confusion_pairs.sort(reverse=True)
for count, true_class, pred_class in confusion_pairs[:5]:
    print(f"  {true_class:12s} → {pred_class:12s}: {count:4d} times")

print("Visualizing model predictions...\n")

# Obtain one test batch
data_iter = iter(test_dataloader)
test_images, test_labels = next(data_iter)

# Make predictions
model.eval()
with torch.no_grad():
    test_images_device = test_images.to(device)
    outputs = model(test_images_device)
    _, predictions = torch.max(outputs, 1)
    predictions = predictions.cpu()

# Visualize first 8 predictions
print("First 8 predictions:")
show_images(test_images[:8], test_labels[:8], predictions[:8])

# Find misclassified examples
incorrect_indices = (predictions != test_labels).nonzero(as_tuple=True)[0]

if len(incorrect_indices) >= 8:
    print("\nExamples of incorrect predictions:")
    error_indices = incorrect_indices[:8]
    show_images(
        test_images[error_indices],
        test_labels[error_indices],
        predictions[error_indices],
    )
else:
    print(f"\nOnly {len(incorrect_indices)} errors in this batch")

import os

os.makedirs("results", exist_ok=True)

def get_activation_maps(model, image, layer_name):
    """
    Extracts activation maps from a specific layer.
    """
    activations = {}

    def hook_fn(module, input, output):
        activations["output"] = output

    # Register hook
    layer = dict([*model.named_modules()])[layer_name]
    hook = layer.register_forward_hook(hook_fn)

    # Forward pass
    model.eval()
    with torch.no_grad():
        _ = model(image.unsqueeze(0).to(device))

    hook.remove()
    return activations["output"].squeeze().cpu()

# Select one test image
test_image, test_label = test_dataset[0]
print(f"Analyzing image of class: {CIFAR10_CLASSES[test_label]}\n")

# Show original image
plt.figure(figsize=(4, 4))
img_display = test_image / 2 + 0.5  # Denormalize
plt.imshow(img_display.permute(1, 2, 0))
plt.title(f"Original Image: {CIFAR10_CLASSES[test_label]}", fontweight="bold")
plt.axis("off")
plt.tight_layout()
plt.show()

# Select and visualize activations from different blocks
layers_to_visualize = {
    "block1": "block1.0",  # First conv of block1
    "block2": "block2.0",  # First conv of block2
    "block3": "block3.0",  # First conv of block3
    "block5": "block5.0",  # First conv of block5
}

for block_name, layer_name in layers_to_visualize.items():
    print(f"Visualizing activations of {block_name}...")

    activations = get_activation_maps(model, test_image, layer_name)

    # Select first 16 filters for visualization
    num_filters = min(16, activations.shape[0])

    fig, axes = plt.subplots(4, 4, figsize=(12, 12))
    fig.suptitle(
        f"Activations of {block_name.upper()} - " f"{CIFAR10_CLASSES[test_label]}",
        fontsize=16,
        fontweight="bold",
    )

    for idx, ax in enumerate(axes.flat):
        if idx < num_filters:
            activation = activations[idx]
            ax.imshow(activation, cmap="viridis")
            ax.set_title(f"Filter {idx}", fontsize=10)
        ax.axis("off")

    plt.tight_layout()
    plt.savefig(f"results/vgg16_activations_{block_name}.png", dpi=300, bbox_inches="tight")
    plt.show()

print("\nInterpretation of block-wise activations:")
print("=" * 70)
print("  Block 1: Detects low-level features")
print("           (edges, corners, simple color variations)")
print("  Block 2-3: Detect mid-level patterns")
print("             (textures, more complex shapes, repetitive patterns)")
print("  Block 4-5: Detect high-level features")
print("             (object parts, combinations of textures and shapes)")