import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# Classes in CIFAR-10
CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck']

# Load the raw (un-normalized) dataset just for visualization
raw_transform = transforms.ToTensor()  # converts PIL image to tensor [0,1]

trainset_raw = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=raw_transform
)
testset_raw = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=raw_transform
)

print(f"Training samples: {len(trainset_raw)}")   # 50000
print(f"Test samples:     {len(testset_raw)}")    # 10000

# Check one sample
img, label = trainset_raw[0]
print(f"Image shape: {img.shape}")     # torch.Size([3, 32, 32])  (C, H, W)
print(f"Image dtype: {img.dtype}")     # torch.float32
print(f"Pixel range: [{img.min():.2f}, {img.max():.2f}]")  # [0.00, 1.00]
print(f"Label: {label} ({CLASSES[label]})")

def show_cifar_samples(dataset, classes, n_per_class=5):
    """Show n_per_class random images from each class."""
    fig, axes = plt.subplots(10, n_per_class, figsize=(n_per_class * 1.5, 16))

    # Collect indices per class
    class_indices = {c: [] for c in range(10)}
    for idx, (_, label) in enumerate(dataset):
        class_indices[label].append(idx)

    for class_id in range(10):
        sample_ids = np.random.choice(class_indices[class_id], n_per_class, replace=False)
        for col, idx in enumerate(sample_ids):
            img, _ = dataset[idx]
            # Convert from (C,H,W) to (H,W,C) for matplotlib
            img_np = img.permute(1, 2, 0).numpy()
            axes[class_id, col].imshow(img_np)
            axes[class_id, col].axis('off')
            if col == 0:
                axes[class_id, col].set_ylabel(classes[class_id], rotation=0,
                                                labelpad=50, fontsize=9)
    plt.suptitle('CIFAR-10: 5 samples from each of 10 classes', y=1.01)
    plt.tight_layout()
    plt.show()

show_cifar_samples(trainset_raw, CLASSES)

import torch
import torchvision
import torchvision.transforms as transforms

# Load raw training data to compute statistics
raw_data = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=False,
    transform=transforms.ToTensor()
)
loader = torch.utils.data.DataLoader(raw_data, batch_size=1000, shuffle=False)

# Compute mean and std per channel across all training images
all_images = torch.cat([x for x, _ in loader])  # (50000, 3, 32, 32)

mean = all_images.mean(dim=(0, 2, 3))  # mean over batch, H, W; one value per channel
std = all_images.std(dim=(0, 2, 3))    # std over batch, H, W

print(f"CIFAR-10 mean per channel: {mean}")
print(f"CIFAR-10 std  per channel: {std}")
# Expected: mean ≈ [0.4914, 0.4822, 0.4465]
#           std  ≈ [0.2470, 0.2435, 0.2616]

# Standard CIFAR-10 normalization values (widely used in published benchmarks)
CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR10_STD  = (0.2470, 0.2435, 0.2616)

# transforms.Normalize(mean, std) applies: output = (input - mean) / std
# per channel. Result: approximately zero mean, unit std.
normalize = transforms.Normalize(mean=CIFAR10_MEAN, std=CIFAR10_STD)

# Basic transform (for validation and test — no random augmentation)
eval_transform = transforms.Compose([
    transforms.ToTensor(),      # PIL image -> float tensor [0,1]
    normalize                   # -> approx zero mean, unit std
])

# Verify normalization
sample, _ = raw_data[0]
normalized = normalize(sample)
print(f"Before normalize: min={sample.min():.3f}, max={sample.max():.3f}")
print(f"After normalize:  min={normalized.min():.3f}, max={normalized.max():.3f}")
print(f"After normalize:  mean≈{normalized.mean():.3f}, std≈{normalized.std():.3f}")

import torch
import torch.nn as nn

class CIFAR10CNN(nn.Module):
    """
    CNN for CIFAR-10 (32x32 RGB, 10 classes).
    Shape trace:
      Input:   (B, 3, 32, 32)
      Block1:  (B, 32, 16, 16)  [conv(3->32,3x3,P=1) + BN + ReLU + Pool(2,2)]
      Block2:  (B, 64, 8, 8)   [conv(32->64,3x3,P=1) + BN + ReLU + Pool(2,2)]
      Block3:  (B, 128, 4, 4)  [conv(64->128,3x3,P=1) + BN + ReLU + Pool(2,2)]
      Flatten: (B, 2048)
      FC1:     (B, 512)
      FC2:     (B, 10)
    """
    def __init__(self, num_classes=10):
        super().__init__()

        self.block1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),  # extra conv in block
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            nn.Dropout2d(0.1)
        )

        self.block2 = nn.Sequential(
            nn.Conv2d(32, 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(2, 2),
            nn.Dropout2d(0.2)
        )

        self.block3 = 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(2, 2),
            nn.Dropout2d(0.3)
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        return self.classifier(x)


# Inspect the model
model = CIFAR10CNN()

# Forward pass to verify shapes
x = torch.randn(4, 3, 32, 32)
print("Shape trace:")
b1 = model.block1(x);    print(f"  After Block 1: {b1.shape}")
b2 = model.block2(b1);   print(f"  After Block 2: {b2.shape}")
b3 = model.block3(b2);   print(f"  After Block 3: {b3.shape}")
out = model.classifier(b3); print(f"  Output:        {out.shape}")

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"\nTotal parameters:     {total_params:,}")
print(f"Trainable parameters: {trainable:,}")

import torchvision.transforms as transforms

CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR10_STD  = (0.2470, 0.2435, 0.2616)

# Training transform — aggressive augmentation
train_transform = transforms.Compose([
    # 1. Random horizontal flip (50% probability)
    transforms.RandomHorizontalFlip(p=0.5),

    # 2. Random crop with padding=4 (shifts image by up to 4 pixels)
    transforms.RandomCrop(32, padding=4),

    # 3. Color jitter (optional but helpful for color robustness)
    transforms.ColorJitter(
        brightness=0.2,   # vary brightness ±20%
        contrast=0.2,     # vary contrast ±20%
        saturation=0.2,   # vary saturation ±20%
        hue=0.1           # vary hue ±10%
    ),

    # 4. Convert to tensor [0, 1]
    transforms.ToTensor(),

    # 5. Normalize to zero mean, unit std
    transforms.Normalize(mean=CIFAR10_MEAN, std=CIFAR10_STD),

    # 6. Random Erasing — erase a random patch (applied after ToTensor)
    transforms.RandomErasing(p=0.5, scale=(0.02, 0.2)),
])

# Validation / Test transform — NO random transforms (deterministic)
eval_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=CIFAR10_MEAN, std=CIFAR10_STD),
])

# Load datasets with the respective transforms
import torchvision
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=train_transform
)
testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=eval_transform
)

print(f"Training samples: {len(trainset)}")  # 50000
print(f"Test samples:     {len(testset)}")   # 10000

# DataLoaders
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True
)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=256, shuffle=False, num_workers=4, pin_memory=True
)

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

CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck']

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

# Model, loss, optimizer, scheduler
model = CIFAR10CNN(num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Cosine annealing: smoothly decays lr from initial to 0 over all epochs
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30)


def train_epoch(model, loader, criterion, optimizer, device):
    """Run one training epoch. Returns (loss, accuracy)."""
    model.train()
    total_loss, correct, total = 0.0, 0, 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += images.size(0)

    return total_loss / total, 100.0 * correct / total


def eval_epoch(model, loader, criterion, device):
    """Evaluate on a data loader. Returns (loss, accuracy)."""
    model.eval()
    total_loss, correct, total = 0.0, 0, 0

    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            total_loss += loss.item() * images.size(0)
            _, predicted = outputs.max(1)
            correct += predicted.eq(labels).sum().item()
            total += images.size(0)

    return total_loss / total, 100.0 * correct / total


# Training loop
NUM_EPOCHS = 30
history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
best_val_acc = 0.0

print(f"\n{'Epoch':>5} {'LR':>8} {'Train Loss':>11} {'Train Acc':>10} {'Val Loss':>10} {'Val Acc':>9}")
print("-" * 60)

for epoch in range(1, NUM_EPOCHS + 1):
    train_loss, train_acc = train_epoch(model, trainloader, criterion, optimizer, device)
    val_loss, val_acc = eval_epoch(model, testloader, criterion, device)
    scheduler.step()

    history['train_loss'].append(train_loss)
    history['train_acc'].append(train_acc)
    history['val_loss'].append(val_loss)
    history['val_acc'].append(val_acc)

    lr = optimizer.param_groups[0]['lr']
    print(f"{epoch:>5} {lr:>8.5f} {train_loss:>11.4f} {train_acc:>9.2f}% {val_loss:>10.4f} {val_acc:>8.2f}%")

    # Save best model
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), 'cifar10_best.pth')
        print(f"          ↑ New best: {best_val_acc:.2f}%")

print(f"\nBest validation accuracy: {best_val_acc:.2f}%")

import torch
import matplotlib.pyplot as plt

# Load the best model
model = CIFAR10CNN().to(device)
model.load_state_dict(torch.load('cifar10_best.pth', map_location=device))
model.eval()

# Register a hook to capture activations after Block 1
activations = {}

def make_hook(name):
    def hook(module, input, output):
        activations[name] = output.detach()
    return hook

# Register hook on the first conv layer and after block1
first_conv = model.block1[0]   # the first nn.Conv2d
handle = first_conv.register_forward_hook(make_hook('conv1_output'))

# Run one test image through the model
testset_vis = torchvision.datasets.CIFAR10(
    root='./data', train=False, transform=eval_transform
)
img, label = testset_vis[42]   # pick an arbitrary image
input_tensor = img.unsqueeze(0).to(device)  # add batch dim: (1, 3, 32, 32)

with torch.no_grad():
    _ = model(input_tensor)

# Remove the hook
handle.remove()

# Visualize the first 32 feature maps from conv1
feat_maps = activations['conv1_output'][0]  # shape: (32, 32, 32)
fig, axes = plt.subplots(4, 8, figsize=(14, 7))
for i, ax in enumerate(axes.flat):
    ax.imshow(feat_maps[i].cpu().numpy(), cmap='RdBu_r')
    ax.axis('off')
    ax.set_title(f'F{i+1}', fontsize=7)
plt.suptitle(f'32 Feature Maps after Conv1 (input: {CLASSES[label]})', y=1.01)
plt.tight_layout()
plt.show()

print(f"Feature map shape: {feat_maps.shape}")  # (32, 32, 32)

import torch
import torch.nn as nn

# BatchNorm2d(32) handles 32-channel feature maps
bn = nn.BatchNorm2d(32)

x = torch.randn(8, 32, 16, 16)  # batch=8, 32 channels, 16x16 spatial
out = bn(x)

print(f"Input  mean (ch 0): {x[:, 0, :, :].mean():.4f}")   # arbitrary
print(f"Output mean (ch 0): {out[:, 0, :, :].mean():.4f}")  # ≈ 0
print(f"Output std  (ch 0): {out[:, 0, :, :].std():.4f}")   # ≈ 1

# The learnable parameters: gamma (scale) and beta (shift)
print(f"BatchNorm parameters: gamma shape={bn.weight.shape}, beta shape={bn.bias.shape}")
# Both have shape (32,) — one per channel