# ============================
# Section 1: General Utilities
# ============================

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
import plotly.graph_objects as go
from IPython.display import clear_output, display

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

gpu_info = !nvidia-smi
gpu_info = '\n'.join(gpu_info)
if gpu_info.find('failed') >= 0:
  print('Not connected to a GPU')
else:
  print(gpu_info)

import plotly.io as pio
import plotly.graph_objects as go

# mont the google drive to save figs
from google.colab import drive
drive.mount('/content/drive')

class ExperimentLogger:
    def __init__(self):
        self.train_losses = []
        self.train_accuracies = []
        self.test_losses = []
        self.test_accuracies = []
        self.lrs = []

    def log_learning_rate(self, lr):
        self.lrs.append(lr)

    def log_train_loss(self, loss):
        self.train_losses.append(loss)

    def log_train_accuracy(self, accuracy):
        self.train_accuracies.append(accuracy)

    def log_test_loss(self, loss):
        self.test_losses.append(loss)

    def log_test_accuracy(self, accuracy):
        self.test_accuracies.append(accuracy)

    def log_train(self, loss, accuracy):
        self.log_train_loss(loss)
        self.log_train_accuracy(accuracy)

    def log_test(self, loss, accuracy):
        self.log_test_loss(loss)
        self.log_test_accuracy(accuracy)

    def live_plot(self, epochs, label):
        clear_output(wait=True)
        fig = go.Figure()

        # Primary y-axis (losses/accuracy)
        fig.add_trace(go.Scatter(y=self.train_losses, mode='lines+markers', name='Training Loss'))
        fig.add_trace(go.Scatter(y=self.train_accuracies, mode='lines+markers', name='Training Accuracy'))
        fig.add_trace(go.Scatter(y=self.test_losses, mode='lines+markers', name='Test Loss'))
        fig.add_trace(go.Scatter(y=self.test_accuracies, mode='lines+markers', name='Test Accuracy'))

        # Secondary y-axis (learning rate)
        if self.lrs:
            fig.add_trace(go.Scatter(y=self.lrs, mode='lines+markers', name='Learning Rate', yaxis='y2'))

        fig.update_layout(
            title=f'{label} - Learning Curves',
            xaxis_title='Epoch',
            yaxis=dict(title='Loss / Accuracy'),
            yaxis2=dict(
                title='Learning Rate',
                overlaying='y',
                side='right',
                showgrid=False
            ),
            width=900,
            height=500,
            legend_title='Metric'
        )

        display(fig)
        
        pio.write_html(fig, file='/content/drive/MyDrive/Colab Notebooks/ResNetPytorch.html', auto_open=False)


    def early_stopping_triggered(self, patience=5, delta=1e-4):
        if len(self.test_accuracies) < patience:
            return False
        recent = self.test_accuracies[-patience:]
        improvement = max(recent) - min(recent)
        return improvement < delta

def show_model_structure(model, name):
    print(f"\n{name} Model Structure:")
    print(model)

# ============================
# Section 2: Dataset Loading
# ============================

from torchvision.transforms import AutoAugment, AutoAugmentPolicy

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    AutoAugment(policy=AutoAugmentPolicy.CIFAR10),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010)),
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

# ============================
# Section 3: CNN Architectures
# ============================


def get_resnet50_for_cifar10():
    model = models.resnet50(pretrained=False)
    model.fc = nn.Linear(model.fc.in_features, 10)
    return model

# Define optimizer and learning rate scheduler
model = get_resnet50_for_cifar10()

# ============================
# Section 4: Model Training & Evaluation
# ============================

def train(model, loader, optimizer, epoch, logger):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

        pred = output.argmax(dim=1)
        correct += (pred == target).sum().item()
        total += target.size(0)

    avg_loss = total_loss / len(loader)
    accuracy = correct / total
    print(f"Train Epoch: {epoch}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.4f}")
    logger.log_train(avg_loss, accuracy)

def test(model, loader, logger):
    model.eval()
    correct = 0
    total = 0
    total_loss = 0
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = F.cross_entropy(output, target)
            total_loss += loss.item()

            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)

    avg_loss = total_loss / len(loader)
    acc = correct / total
    print(f"Test Accuracy: {acc:.4f}, Loss: {avg_loss:.4f}")
    logger.log_test(avg_loss, acc)

# ============================
# Section 5: Run Experiments
# ============================

import copy  # for deep copying model state_dict

def run_experiment(name, model, train_dataset, test_dataset, epochs,
                   optimizer_class=optim.SGD, optimizer_kwargs=None,
                   optimizer=None, lr_scheduler=None, 
                   batch_size=64, val_batch_size=64,
                   patience=None, delta=1e-4, 
                   checkpoint_path=None  # path to save the best model
                   ):
    
    best_model_state = None
    model = model.to(device)
    show_model_structure(model, name)

    logger = ExperimentLogger()

    # Use provided optimizer or construct one
    if optimizer is None:
        optimizer_kwargs = optimizer_kwargs or {"lr": 0.1, "momentum": 0.9}
        optimizer = optimizer_class(model.parameters(), **optimizer_kwargs)

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

    best_accuracy = 0
    epochs_no_improve = 0

    for epoch in range(1, epochs + 1):
        train(model, train_loader, optimizer, epoch, logger)
        test(model, test_loader, logger)

        current_lr = optimizer.param_groups[0]['lr']
        logger.log_learning_rate(current_lr)

        if lr_scheduler:
            lr_scheduler.step()

        if logger.test_accuracies[-1] > best_accuracy + delta:
            best_accuracy = logger.test_accuracies[-1]
            best_model_state = copy.deepcopy(model.state_dict())  # 💾 store best model
            epochs_no_improve = 0
        else:
            epochs_no_improve += 1
            if patience and epochs_no_improve >= patience:
                print(f"Early stopping at epoch {epoch} (no improvement for {patience} epochs)")
                break

        logger.live_plot(epoch, name)

    if checkpoint_path and best_model_state:
        torch.save(best_model_state, checkpoint_path)
        print(f"✅ Saved best model to: {checkpoint_path}")

    return model, logger