# ==============================
# 0. Install & Imports
# ==============================
!pip install --quiet gdown
from tqdm import tqdm, trange
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split, TensorDataset
from torchvision import datasets, transforms, models
import pandas as pd
import gdown

# ==============================
# 1. Hyperparameters & Device
# ==============================
BATCH_SIZE = 64
LR = 1e-3
EPOCHS = 10
N_CLASSES = 10
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_DIR = "./CIFAR10_data"
os.makedirs(DATA_DIR, exist_ok=True)

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

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

# ==============================
# 3. Load Train/Validation (Public CIFAR-10)
# ==============================
full_train_dataset = datasets.CIFAR10(root=DATA_DIR, train=True, download=True, transform=transform_train)
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)

# ==============================
# 4. Download Test Dataset
# ==============================
# Students download this version only (images, no labels)
FILE_ID = "1cXUWjw-WiJqXb5gfjtqGFDMDDOGi2N56" 
DEST_PATH = "cifar10_test.pt"
gdown.download(f"https://drive.google.com/uc?id={FILE_ID}", DEST_PATH, quiet=False)

# Load images only
test_images = torch.load(DEST_PATH)  # shape [N, 3, 32, 32]

# Dummy labels for DataLoader
test_labels = torch.zeros(len(test_images), dtype=torch.long)
test_dataset = TensorDataset(test_images, test_labels)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

print(f"Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")

# ==============================
# 5. Define Models
# ==============================

# 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)
)

models_dict = {
    "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, show_progress=False):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    y_pred = []
    iterator = loader
    if show_progress:
        iterator = tqdm(loader, desc="Evaluating", unit="batch")
    with torch.no_grad():
        for x, y in iterator:
            x, y = x.to(DEVICE), y.to(DEVICE)
            out = model(x)
            if criterion is not None:
                total_loss += criterion(out, y).item() * x.size(0)
            preds = out.argmax(1)
            y_pred.extend(preds.cpu().numpy())
            correct += (preds == y).sum().item()
            total += y.size(0)
    avg_loss = None if criterion is None else total_loss / total
    acc = correct / total
    return avg_loss, acc, y_pred

# ==============================
# 7. Training Loop + Save Test Predictions
# ==============================
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} ...")

    # Epoch loop with progress bar
    for epoch in trange(1, EPOCHS + 1, desc="Epochs", unit="epoch"):
        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)
        val_loss, val_acc, _ = evaluate(model, val_loader, criterion)
        tqdm.write(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}")

    # Evaluate on test set with progress bar
    _, _, test_preds = evaluate(model, test_loader, show_progress=True)
    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}")