# ==============================
# 0. Mount Google Drive (optional)
# ==============================
from google.colab import drive
import os

drive.mount('/content/drive')

DATA_DIR = '/content/drive/MyDrive/CIFAR10_data'
os.makedirs(DATA_DIR, exist_ok=True)
print(f"Dataset will be stored in: {DATA_DIR}")

# ==============================
# 1. Imports
# ==============================
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms, models
import pandas as pd

# ==============================
# 2. Hyperparameters
# ==============================
BATCH_SIZE = 64
LR = 1e-3
EPOCHS = 10
N_CLASSES = 10
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ==============================
# 3. Data Transforms
# ==============================
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor()
])

transform_test = transforms.Compose([
    transforms.ToTensor()
])

# ==============================
# 4. Dataset & DataLoaders
# ==============================
# Use the same test dataset for all students
full_train_dataset = datasets.CIFAR10(root=DATA_DIR, train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10(root=DATA_DIR, train=False, download=True, transform=transform_test)

# Split train into train/val
train_size = int(0.8 * len(full_train_dataset))
val_size = len(full_train_dataset) - train_size
train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])

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

print("DataLoaders ready: train, val, test")

# ==============================
# 5. Define Models
# ==============================
class MyCNN(nn.Module):
    def __init__(self, num_classes=N_CLASSES):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64*8*8, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.conv(x)
        x = self.fc(x)
        return x

# VGG16 Transfer Learning
vgg_model = models.vgg16(weights=models.VGG16_Weights.DEFAULT)
for param in vgg_model.features.parameters():
    param.requires_grad = False
num_features = vgg_model.classifier[0].in_features
vgg_model.classifier = nn.Sequential(
    nn.Linear(num_features, 512),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(512, N_CLASSES)
)

# Move models to device
models_dict = {
    "CustomCNN": MyCNN().to(DEVICE),
    "VGG16": vgg_model.to(DEVICE)
}

# ==============================
# 6. Training & Evaluation Functions
# ==============================
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for x, y in loader:
        x, y = x.to(DEVICE), y.to(DEVICE)
        optimizer.zero_grad()
        out = model(x)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * x.size(0)
        correct += (out.argmax(1) == y).sum().item()
        total += y.size(0)
    return total_loss / total, correct / total

def evaluate(model, loader, criterion=None):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    y_true, y_pred = [], []
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            out = model(x)
            if criterion is not None:
                loss = criterion(out, y)
                total_loss += loss.item() * x.size(0)
            preds = out.argmax(1)
            correct += (preds == y).sum().item()
            total += y.size(0)
            y_true.extend(y.cpu().numpy())
            y_pred.extend(preds.cpu().numpy())

    avg_loss = total_loss / total if criterion is not None else None
    acc = correct / total
    return avg_loss, acc, y_true, y_pred

# ==============================
# 7. Training Loop
# ==============================
criterion = nn.CrossEntropyLoss()

for name, model in models_dict.items():
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)
    print(f"\nTraining {name} ...")
    for epoch in range(1, EPOCHS + 1):
        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)
        val_loss, val_acc, _, _ = evaluate(model, val_loader, criterion)
        print(f"Epoch [{epoch}/{EPOCHS}] | "
              f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.3f} | "
              f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.3f}")

    # ==============================
    # Save Test Predictions to CSV for Leaderboard
    # ==============================
    _, _, _, test_preds = evaluate(model, test_loader)
    submission = pd.DataFrame({
        "Id": list(range(len(test_preds))),
        "Predicted": test_preds
    })
    submission_file = f"{name}_CIFAR10_test_predictions.csv"
    submission.to_csv(submission_file, index=False)
    print(f"Test predictions saved to {submission_file}")