# -- Code Cell --
import time
import math
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Subset, random_split, Dataset
from torchvision import datasets, transforms
import os
import pandas as pd
from PIL import Image

# -- Code Cell --
class EmojiDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None, is_test=False):
        self.data = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform
        self.is_test = is_test

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        
        img_path = row["Path"]
        image = Image.open(img_path).convert("RGB")

        if self.transform:
            image = self.transform(image)

        if self.is_test:
            return image, row["SampleID"]
        else:
            label = row["Label"]
            return image, label


# -- Code Cell --
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

train_dataset = EmojiDataset(
    csv_file="train.csv",
    root_dir="train",
    transform=transform,
    is_test=False
)

test_dataset = EmojiDataset(
    csv_file="test.csv",
    root_dir="test",
    transform=transform,
    is_test=True
)


labels = train_dataset.data["Label"].astype(int)
NUM_CLASSES = int(labels.max()) + 1

# -- Code Cell --
train_subset_size = len(train_dataset)
val_size = int(0.15 * train_subset_size)

train_dataset, val_dataset = random_split(
    train_dataset,
    [train_subset_size - val_size, val_size]
)


# -- Code Cell --
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# -- Code Cell --
def train_one_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for x, y in dataloader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        logits = model(x)
        loss = criterion(logits, y)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * x.size(0)
        _, preds = torch.max(logits, 1)
        correct += (preds == y).sum().item()
        total += y.size(0)

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def evaluate(model, dataloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for x, y in dataloader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            loss = criterion(logits, y)

            running_loss += loss.item() * x.size(0)
            _, preds = torch.max(logits, 1)
            correct += (preds == y).sum().item()
            total += y.size(0)

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def train_model(model, train_loader, val_loader, optimizer, criterion, device, epochs=3):
    model.to(device)
    history = {
        "train_loss": [],
        "train_acc": [],
        "val_loss": [],
        "val_acc": [],
        "epoch_time": []
    }

    for epoch in range(1, epochs + 1):
        start = time.time()
        train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion, device)
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
        end = time.time()

        history["train_loss"].append(train_loss)
        history["train_acc"].append(train_acc)
        history["val_loss"].append(val_loss)
        history["val_acc"].append(val_acc)
        history["epoch_time"].append(end - start)

        print(
            f"Epoch {epoch}/{epochs} | "
            f"Train loss: {train_loss:.4f}, acc: {train_acc:.4f} | "
            f"Val loss: {val_loss:.4f}, acc: {val_acc:.4f} | "
            f"Time: {end - start:.1f}s"
        )

    return history


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# We'll collect all model results in a list of dicts for later summary.
model_results = []

# -- Code Cell --
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -- Code Cell --
# ============================================
# 5. Simple CNN (Conv → ReLU → Pool → Conv → ReLU → Pool → Flatten → Linear → ReLU → Linear)
# ============================================

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=NUM_CLASSES):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),  
            nn.ReLU(),
            nn.MaxPool2d(2), 

            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),  
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 56 * 56, 128), 
            nn.ReLU(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        return self.classifier(x)


cnn = SimpleCNN()
print("SimpleCNN parameters:", count_parameters(cnn))

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(cnn.parameters(), lr=1e-3)

cnn_history = train_model(
    cnn,
    train_loader,
    val_loader,
    optimizer,
    criterion,
    device,
    epochs=3
)

cnn_val_acc = cnn_history["val_acc"][-1]
cnn_epoch_time = np.mean(cnn_history["epoch_time"])

model_results.append({
    "name": "SimpleCNN",
    "params": count_parameters(cnn),
    "val_acc": cnn_val_acc,
    "epoch_time": cnn_epoch_time
})


# ----- Intermediary Performance Comparison B: MLP vs SimpleCNN -----
print("\n=== Intermediary Comparison B: MLP vs SimpleCNN ===")
df = pd.DataFrame(model_results)
print(df)