import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

# Set seed for reproducibility
torch.manual_seed(42)

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

BATCH_SIZE = 128

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")

# Look at some examples
examples = iter(train_loader)
images, labels = next(examples)

fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for i, ax in enumerate(axes.flat):
    ax.imshow(images[i].squeeze(), cmap='gray')
    ax.set_title(f'Label: {labels[i]}')
    ax.axis('off')
plt.tight_layout()
plt.show()

print(f"Image shape: {images[0].shape}")
print(f"Min pixel value: {images[0].min():.3f}")
print(f"Max pixel value: {images[0].max():.3f}")

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear = nn.Linear(784, 10)  # 28x28 = 784
    
    def forward(self, x):
        x = self.flatten(x)
        return self.linear(x)

# Quick inspection
model = LinearModel()
print(model)
print(f"\nNumber of parameters: {sum(p.numel() for p in model.parameters()):,}")

class MLP(nn.Module):
    def __init__(self, hidden_size=128):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(784, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, 10)
    
    def forward(self, x):
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        return self.fc2(x)

model = MLP()
print(model)
print(f"\nNumber of parameters: {sum(p.numel() for p in model.parameters()):,}")

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.relu = nn.ReLU()
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))  # 28x28 -> 14x14
        x = self.pool(self.relu(self.conv2(x)))  # 14x14 -> 7x7
        x = x.view(-1, 32 * 7 * 7)
        x = self.relu(self.fc1(x))
        return self.fc2(x)

model = SimpleCNN()
print(model)
print(f"\nNumber of parameters: {sum(p.numel() for p in model.parameters()):,}")

def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    
    return total_loss / len(loader), 100. * correct / total

def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 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()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    return total_loss / len(loader), 100. * correct / total

EPOCHS = 5
LEARNING_RATE = 0.001

models = {
    'Linear': LinearModel(),
    'MLP': MLP(),
    'CNN': SimpleCNN()
}

results = {}

for name, model in models.items():
    print(f"\n{'='*60}")
    print(f"Training {name}")
    print(f"{'='*60}")
    
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    
    train_accs = []
    test_accs = []
    
    for epoch in range(EPOCHS):
        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
        test_loss, test_acc = evaluate(model, test_loader, criterion, device)
        
        train_accs.append(train_acc)
        test_accs.append(test_acc)
        
        print(f"Epoch {epoch+1}/{EPOCHS} | "
              f"Train: {train_acc:.2f}% | Test: {test_acc:.2f}%")
    
    results[name] = {
        'train_acc': train_accs,
        'test_acc': test_accs,
        'final': test_acc
    }

print("\nDone!")

# Plot comparison
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

epochs = range(1, EPOCHS + 1)

# Test accuracy over time
ax = axes[0]
for name, data in results.items():
    ax.plot(epochs, data['test_acc'], marker='o', label=name, linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Test Accuracy (%)')
ax.set_title('Test Accuracy vs Epoch')
ax.legend()
ax.grid(True, alpha=0.3)

# Final comparison
ax = axes[1]
names = list(results.keys())
final_accs = [results[n]['final'] for n in names]
bars = ax.bar(names, final_accs, alpha=0.7, edgecolor='black')
for bar, acc in zip(bars, final_accs):
    height = bar.get_height()
    ax.text(bar.get_x() + bar.get_width()/2., height,
            f'{acc:.2f}%', ha='center', va='bottom', fontweight='bold')
ax.set_ylabel('Test Accuracy (%)')
ax.set_title('Final Test Accuracy')
ax.set_ylim([90, 100])
ax.grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.show()

# Get more diverse samples
fig, axes = plt.subplots(10, 10, figsize=(12, 12))

data_iter = iter(train_loader)
images, labels = next(data_iter)

for i in range(10):
    for j in range(10):
        idx = i * 10 + j
        img = images[idx].squeeze() * 0.3081 + 0.1307  # Denormalize
        axes[i, j].imshow(img, cmap='gray', vmin=0, vmax=1)
        axes[i, j].set_title(f'{labels[idx].item()}', fontsize=8)
        axes[i, j].axis('off')

plt.suptitle('100 Random MNIST Samples', fontsize=16, fontweight='bold')
plt.tight_layout()
plt.show()

print("\nKey observations:")
print("✓ Digits are centered and roughly same scale")
print("✓ Background is consistent (dark)")
print("✓ Foreground is consistent (light)")
print("✓ Limited style variation (clean handwriting)")
print("\n→ This structure makes linear separation easier!")

from sklearn.decomposition import PCA
from sklearn.manifold import TSNE

# Collect subset for visualization (t-SNE is slow)
n_samples = 2000
X_viz = []
y_viz = []

for images, labels in train_loader:
    X_viz.append(images.view(images.size(0), -1))
    y_viz.append(labels)
    if len(X_viz) * BATCH_SIZE >= n_samples:
        break

X_viz = torch.cat(X_viz, dim=0)[:n_samples].numpy()
y_viz = torch.cat(y_viz, dim=0)[:n_samples].numpy()

print(f"Projecting {n_samples} samples from 784D to 2D...")

# PCA (linear)
pca = PCA(n_components=2)
X_pca = pca.fit_transform(X_viz)
print(f"PCA explained variance: {pca.explained_variance_ratio_.sum():.1%}")

# t-SNE (nonlinear, slow)
print("Running t-SNE (30-60 seconds)...")
tsne = TSNE(n_components=2, random_state=42, perplexity=30)
X_tsne = tsne.fit_transform(X_viz)
print("Done!")

# Plot both projections
fig, axes = plt.subplots(1, 2, figsize=(16, 6))
colors = plt.cm.tab10(np.linspace(0, 1, 10))

# PCA
ax = axes[0]
for digit in range(10):
    mask = y_viz == digit
    ax.scatter(X_pca[mask, 0], X_pca[mask, 1],
              c=[colors[digit]], label=str(digit), alpha=0.6, s=20)
ax.set_xlabel('PC1', fontsize=12)
ax.set_ylabel('PC2', fontsize=12)
ax.set_title(f'PCA (Linear) - {pca.explained_variance_ratio_.sum():.1%} variance',
             fontsize=14, fontweight='bold')
ax.legend(title='Digit', bbox_to_anchor=(1.05, 1), loc='upper left')
ax.grid(True, alpha=0.3)

# t-SNE
ax = axes[1]
for digit in range(10):
    mask = y_viz == digit
    ax.scatter(X_tsne[mask, 0], X_tsne[mask, 1],
              c=[colors[digit]], label=str(digit), alpha=0.6, s=20)
ax.set_xlabel('t-SNE 1', fontsize=12)
ax.set_ylabel('t-SNE 2', fontsize=12)
ax.set_title('t-SNE (Nonlinear) - Better separation',
             fontsize=14, fontweight='bold')
ax.legend(title='Digit', bbox_to_anchor=(1.05, 1), loc='upper left')
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print("\n🔍 Analysis:")
print("- PCA (linear): Classes show SOME separation even with just 2 components!")
print("  → First 2 PCs capture ~17% of variance but reveal class structure")
print("  → This is why linear classification works reasonably well")
print("\n- t-SNE: Much better separation, but this is nonlinear")
print("  → Shows the 'best case' if we could use perfect nonlinear features")

# Get the trained linear model from earlier
# (If you haven't saved it, retrain quickly)
linear_model_viz = models['Linear']  # Should still be in memory

# Extract learned weights
weights = linear_model_viz.linear.weight.data.cpu().numpy()  # Shape: (10, 784)

# Visualize each class's weight vector as 28x28 image
fig, axes = plt.subplots(2, 5, figsize=(15, 6))

for i, ax in enumerate(axes.flat):
    weight_img = weights[i].reshape(28, 28)
    vmax = np.abs(weight_img).max()
    
    im = ax.imshow(weight_img, cmap='RdBu', vmin=-vmax, vmax=vmax)
    ax.set_title(f'Digit {i}', fontsize=12, fontweight='bold')
    ax.axis('off')

plt.suptitle('What the Linear Model Learned\n(Red = positive weight, Blue = negative weight)',
             fontsize=14, fontweight='bold')
plt.colorbar(im, ax=axes.ravel().tolist(), label='Weight', fraction=0.046, pad=0.04)
plt.tight_layout()
plt.show()

print("\n Interpretation:")
print("- Each digit learns a 'template' in pixel space")
print("- Red pixels: presence INCREASES score for this digit")
print("- Blue pixels: presence DECREASES score for this digit")
print("\nClassification = argmax(W·x)")
print("→ Which template matches the input best?")
print("\n️  This is literally TEMPLATE MATCHING in pixel space!")

# Simple linear model for 2D input
class LinearModel2D(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 10)
    def forward(self, x):
        return self.linear(x)

# Train on PCA projection
X_train_tensor = torch.FloatTensor(X_pca)
y_train_tensor = torch.LongTensor(y_viz)

model_2d = LinearModel2D().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model_2d.parameters(), lr=0.01)

print("Training linear model on 2D PCA projection...")
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model_2d(X_train_tensor.to(device))
    loss = criterion(outputs, y_train_tensor.to(device))
    loss.backward()
    optimizer.step()

# Check accuracy
with torch.no_grad():
    _, predicted = model_2d(X_train_tensor.to(device)).max(1)
    acc_2d = (predicted.cpu() == y_train_tensor).float().mean() * 100

print(f"Accuracy on 2D projection: {acc_2d:.2f}%")
print(f"Accuracy on full 784D: ~92%")
print(f"\n→ Higher dimensions allow MUCH better separation!")

# Create decision boundary visualization
h = 0.5  # Mesh step size
x_min, x_max = X_pca[:, 0].min() - 1, X_pca[:, 0].max() + 1
y_min, y_max = X_pca[:, 1].min() - 1, X_pca[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))

# Predict on mesh
mesh_points = torch.FloatTensor(np.c_[xx.ravel(), yy.ravel()]).to(device)
with torch.no_grad():
    Z = model_2d(mesh_points).cpu().numpy()
    Z = np.argmax(Z, axis=1).reshape(xx.shape)

# Plot
plt.figure(figsize=(12, 10))
plt.contourf(xx, yy, Z, alpha=0.3, levels=np.arange(11) - 0.5, cmap='tab10')

# Overlay data points
for digit in range(10):
    mask = y_viz == digit
    plt.scatter(X_pca[mask, 0], X_pca[mask, 1],
               c=[colors[digit]], label=str(digit), alpha=0.7, s=30,
               edgecolors='k', linewidth=0.5)

plt.xlabel('PC1', fontsize=12)
plt.ylabel('PC2', fontsize=12)
plt.title(f'Linear Decision Boundaries in PCA Space\n(Colored regions = predicted class, Accuracy = {acc_2d:.1f}%)',
         fontsize=14, fontweight='bold')
plt.legend(title='Digit', bbox_to_anchor=(1.05, 1), loc='upper left')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

print("\n📐 Key insights:")
print("1. Decision boundaries are LINEAR (straight lines in 2D, hyperplanes in higher D)")
print("2. Even with just 2D, we get SOME separation")
print("3. Look at the center - massive overlap of all classes")
print("4. In 784D, these classes separate much better (92% vs 36%)")
print("5. The 'curse of dimensionality' actually helps us here!")

# Collect misclassifications
linear_model_viz.eval()
misclassified = []

with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = linear_model_viz(images)
        _, predicted = outputs.max(1)
        
        # Find errors
        wrong_idx = (predicted != labels).nonzero(as_tuple=True)[0]
        
        for idx in wrong_idx:
            if len(misclassified) < 20:
                misclassified.append({
                    'image': images[idx].cpu(),
                    'true': labels[idx].cpu().item(),
                    'pred': predicted[idx].cpu().item()
                })
        
        if len(misclassified) >= 20:
            break

# Visualize failures
fig, axes = plt.subplots(4, 5, figsize=(15, 12))

for i, ax in enumerate(axes.flat):
    example = misclassified[i]
    img = example['image'].squeeze() * 0.3081 + 0.1307  # Denormalize
    
    ax.imshow(img, cmap='gray')
    ax.set_title(f"True: {example['true']}, Pred: {example['pred']}",
                fontsize=10, color='red', fontweight='bold')
    ax.axis('off')

plt.suptitle('Linear Model Failures\n(Where template matching breaks down)',
             fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

print("\n Common failure modes:")
print("- Ambiguous handwriting (could be multiple digits)")
print("- Similar digits: 4/9, 3/5/8, 1/7, 6/0")
print("- Unusual styles (too cursive, too geometric)")
print("- Poor centering or scaling")
print("\n These are exactly where NONLINEAR models can help!")
print("   → MLP can learn more complex decision boundaries")
print("   → CNN can learn spatial features (not just templates)")

# Layer 1
conv1 = Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1)
# Input:  (batch, 1, 28, 28)
# Output: (batch, 16, 28, 28)  [16 feature maps]
# Parameters: 16 filters × (1 input channel × 3 × 3) + 16 biases = 160

pool1 = MaxPool2d(kernel_size=2, stride=2)
# Output: (batch, 16, 14, 14)  [spatial downsampling]

# Layer 2
conv2 = Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
# Input:  (batch, 16, 14, 14)
# Output: (batch, 32, 14, 14)  [32 feature maps]
# Parameters: 32 filters × (16 input channels × 3 × 3) + 32 biases = 4,640

pool2 = MaxPool2d(kernel_size=2, stride=2)
# Output: (batch, 32, 7, 7)

# Fully connected
fc1 = Linear(32*7*7, 128)  # Flatten spatial dimensions
fc2 = Linear(128, 10)      # Final classification

# Extract first layer filters
cnn_model_viz = models['CNN']
filters_layer1 = cnn_model_viz.conv1.weight.data.cpu().numpy()

print(f"Filter tensor shape: {filters_layer1.shape}")
print(f"Interpretation: (out_channels={filters_layer1.shape[0]}, "
      f"in_channels={filters_layer1.shape[1]}, "
      f"height={filters_layer1.shape[2]}, "
      f"width={filters_layer1.shape[3]})")

# For visualization, we extract each 3x3 filter
# Since in_channels=1, we can squeeze that dimension
# filters_layer1[i, 0, :, :] gives the i-th 3x3 filter

# Visualize all 16 filters
fig, axes = plt.subplots(4, 4, figsize=(10, 10))

for i, ax in enumerate(axes.flat):
    # Extract i-th filter (shape: 3x3)
    filt = filters_layer1[i, 0, :, :]
    
    # Normalize for visualization (center at 0)
    vmax = np.abs(filt).max()
    
    # Plot with diverging colormap
    im = ax.imshow(filt, cmap='RdBu', vmin=-vmax, vmax=vmax, interpolation='nearest')
    ax.set_title(f'Filter {i}', fontsize=10)
    ax.axis('off')

plt.suptitle('First Conv Layer: 16 Learned 3x3 Filters',
             fontsize=14, fontweight='bold')
plt.colorbar(im, ax=axes.ravel().tolist(), fraction=0.046, pad=0.04, label='Weight value')
plt.tight_layout()
plt.show()

# Get a test image
test_iter = iter(test_loader)
test_images, test_labels = next(test_iter)

# Select example
example_idx = 0
example_img = test_images[example_idx:example_idx+1].to(device)
example_label = test_labels[example_idx].item()

print(f"Input shape: {example_img.shape}")
print(f"True label: {example_label}")

# Forward pass through first conv layer
cnn_model_viz.eval()
with torch.no_grad():
    # Step 1: Convolution
    conv1_out = cnn_model_viz.conv1(example_img)
    print(f"After Conv2d: {conv1_out.shape}")
    
    # Step 2: ReLU activation
    conv1_activated = cnn_model_viz.relu(conv1_out)
    print(f"After ReLU: {conv1_activated.shape}")

# Extract activations (remove batch dimension)
activations_layer1 = conv1_activated.cpu().numpy()[0]  # Shape: (16, 28, 28)
print(f"\nActivations shape: {activations_layer1.shape}")
print(f"Interpretation: 16 feature maps, each 28x28")

# Visualize input and all activations
fig = plt.figure(figsize=(16, 10))

# Show original image
ax = plt.subplot(4, 5, 1)
img_display = example_img.cpu().numpy()[0, 0] * 0.3081 + 0.1307  # Denormalize
ax.imshow(img_display, cmap='gray')
ax.set_title(f'Input: {example_label}', fontsize=14, fontweight='bold')
ax.axis('off')

# Show 16 filter activations
for i in range(16):
    ax = plt.subplot(4, 5, i+2)
    ax.imshow(activations_layer1[i], cmap='viridis')
    ax.set_title(f'Filter {i}', fontsize=9)
    ax.axis('off')

plt.suptitle(f'First Layer Feature Maps for Digit "{example_label}"\n(Brighter = stronger activation)',
             fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

# Forward through conv1 + pool, then conv2
with torch.no_grad():
    # Layer 1
    x = cnn_model_viz.conv1(example_img)
    x = cnn_model_viz.relu(x)
    x = cnn_model_viz.pool(x)
    print(f"After pool1: {x.shape}")
    
    # Layer 2
    x = cnn_model_viz.conv2(x)
    conv2_activated = cnn_model_viz.relu(x)
    print(f"After conv2+ReLU: {conv2_activated.shape}")

# Extract activations
activations_layer2 = conv2_activated.cpu().numpy()[0]  # Shape: (32, 14, 14)
print(f"\nLayer 2 activations: {activations_layer2.shape}")
print(f"Interpretation: 32 feature maps, each 14x14")

# Visualize input and first 16 of 32 second-layer activations
fig = plt.figure(figsize=(16, 8))

# Show input
ax = plt.subplot(4, 5, 1)
ax.imshow(img_display, cmap='gray')
ax.set_title('Input', fontsize=12, fontweight='bold')
ax.axis('off')

# Show 16 second-layer activations (out of 32 total)
for i in range(16):
    ax = plt.subplot(4, 5, i+2)
    ax.imshow(activations_layer2[i], cmap='viridis')
    ax.set_title(f'L2-{i}', fontsize=9)
    ax.axis('off')

plt.suptitle('Second Conv Layer Activations (14x14 feature maps)',
             fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

# Get linear model's weight vector for this digit
linear_model_viz = models['Linear']
linear_weights_all = linear_model_viz.linear.weight.data.cpu().numpy()
linear_template = linear_weights_all[example_label].reshape(28, 28)

# Create comparison visualization
fig, axes = plt.subplots(2, 5, figsize=(15, 6))

# Top row: Linear template + 4 CNN filters
ax = axes[0, 0]
vmax = np.abs(linear_template).max()
ax.imshow(linear_template, cmap='RdBu', vmin=-vmax, vmax=vmax)
ax.set_title(f'Linear\nTemplate for "{example_label}"\n(28x28)', fontsize=10, fontweight='bold')
ax.axis('off')

for i in range(4):
    ax = axes[0, i+1]
    filt = filters_layer1[i, 0, :, :]
    vmax_filt = np.abs(filt).max()
    ax.imshow(filt, cmap='RdBu', vmin=-vmax_filt, vmax=vmax_filt, interpolation='nearest')
    ax.set_title(f'CNN\nFilter {i}\n(3x3)', fontsize=10, fontweight='bold')
    ax.axis('off')

# Bottom row: Input image + 4 activation maps
ax = axes[1, 0]
ax.imshow(img_display, cmap='gray')
ax.set_title('Input\nDigit', fontsize=10, fontweight='bold')
ax.axis('off')

for i in range(4):
    ax = axes[1, i+1]
    ax.imshow(activations_layer1[i], cmap='viridis')
    ax.set_title(f'CNN\nActivation {i}\n(28x28)', fontsize=10, fontweight='bold')
    ax.axis('off')

plt.suptitle('Linear (Global Template) vs CNN (Local Features + Combination)',
             fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()