# -- Markdown Cell --
# # Finetune a ResNet image classifier

# -- Markdown Cell --
# ## Load data

# -- Code Cell --
!curl -L https://storage.googleapis.com/aiolympiadmy/ioai-2025-tsp/finetuning-resnet.zip -o data.zip

# -- Code Cell --
import zipfile

with zipfile.ZipFile('data.zip', 'r') as zip_ref:
    zip_ref.extractall()

# -- Markdown Cell --
# ## Establishing a baseline

# -- Markdown Cell --
# Load an ImageNet pre-trained ResNet34, and check it's performance on the images in `data/test`.  
# 
# Use accuracy, precision and recall as your metrics for performance.

# -- Code Cell --
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import random
from torchvision import datasets, transforms, models
from torchvision.models import ResNet34_Weights
from torch.utils.data import DataLoader, Subset, ConcatDataset, Sampler
from tqdm import tqdm
from sklearn.metrics import accuracy_score, precision_score, recall_score

# -- Markdown Cell --
# ### Utilities

# -- Code Cell --
import os
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from collections import Counter
 
def count_images(data_dir):
    counts = {}
    for split in ['train', 'test']:
        split_path = os.path.join(data_dir, split)
        counts[split] = {}
       
        for cls in os.listdir(split_path):
            cls_path = os.path.join(split_path, cls)
            counts[split][cls] = len(os.listdir(cls_path))
   
    return counts
 
data_dir = "./data"
counts = count_images(data_dir)
 
print(counts)

# -- Code Cell --
def plot_class_distribution(counts):
    for split in counts:
        classes = list(counts[split].keys())
        values = list(counts[split].values())

        print(classes, values)

        plt.figure()
        plt.bar(classes, values)
        plt.title(f"{split.capitalize()} Class Distribution")
        plt.xlabel("Class")
        plt.ylabel("Number of Images")
        plt.show()

plot_class_distribution(counts)


# -- Code Cell --
import random

def show_samples(data_dir, split='train', n=5):
    plt.figure(figsize=(10, 5))
    
    i = 1
    for cls in ['chihuahua', 'muffin']:
        cls_path = os.path.join(data_dir, split, cls)
        images = random.sample(os.listdir(cls_path), n)
        
        for img_name in images:
            img_path = os.path.join(cls_path, img_name)
            img = Image.open(img_path)
            
            plt.subplot(2, n, i)
            plt.imshow(img)
            plt.title(cls)
            plt.axis('off')
            i += 1
    
    plt.tight_layout()
    plt.show()

show_samples(data_dir, "train", 7)

# -- Code Cell --
class FixedOrderSampler(Sampler):
    __slots = ("indices", )
    
    def __init__(self, indices):
        self.indices = indices

    def __iter__(self):
        return iter(self.indices)

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

# -- Code Cell --
from torchvision.transforms import AugMix
from sklearn.model_selection import StratifiedKFold
from torch.utils.data import Subset
from torchvision.transforms import v2

# -- Code Cell --
class BinaryImageClassifier:
    __slots__ = ("__lr", "__l2_penalty", "pretrained", "device", "criterion", "train_loss_history", "test_loss_history",
                 "train_metrics_history", "test_results", "model", "optimizer", "train_loader", "test_loader", "train_dir")

    def __init__(self, *, lr, l2_penalty=0, pretrained=True, train_dir="./data/train", test_dir="./data/test"):
        self.__lr = lr
        self.__l2_penalty = l2_penalty
        self.pretrained = pretrained
        self.train_dir = train_dir
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.criterion = nn.BCELoss()  # Binary Cross Entropy Loss
        self.train_loss_history = []
        self.test_loss_history = []
        self.train_metrics_history = []
        self.test_results = {}

        self.reload_model()

        self.train_loader = self.load_data(train_dir, is_train=True)
        self.test_loader = self.load_data(test_dir, is_train=False)


    @property
    def lr(self):
        return self.__lr

    @lr.setter
    def lr(self, val):
        self.__lr = val
        self.__reload_optimizer()


    @property
    def l2_penalty(self):
        return self.__l2_penalty

    @l2_penalty.setter
    def l2_penalty(self, val):
        self.__l2_penalty = val
        self.__reload_optimizer()


    def __reload_optimizer(self):
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.l2_penalty)

    def fit_kfold(self, n_epochs, n_splits=5):
        dataset = self.train_loader.dataset  
        X = np.arange(len(dataset))
        y = np.array(dataset.targets)
        skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        results = []
        for fold, (train_idx, val_idx) in enumerate(skf.split(X, y), 1):
            print(f"Fold {fold}")
            self.train_loader = DataLoader(Subset(dataset, train_idx),batch_size=64,shuffle=True)
            self.test_loader = DataLoader(Subset(dataset, val_idx),batch_size=64,shuffle=False)
            self.fit(n_epochs=n_epochs, reset=True)
            res = self.predict(show_progress=False)
            results.append(res["metrics"])
            print(res["metrics"])

    def load_data(self, dir_, start=0, stop=None, step=1, is_train=True):
        self.__set_seed()

        if is_train:
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.Lambda(lambda img: img.convert("RGB")),
                AugMix(),
                transforms.ToTensor()
            ])
        else:
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.Lambda(lambda img: img.convert("RGB")),
                transforms.ToTensor()
            ])
        dataset = datasets.ImageFolder(root=dir_, transform=transform)

        # Shuffle the dataset
        indices = list(range(len(dataset)))
        random.shuffle(indices)
        indices = indices[start:stop:step]

        loader = DataLoader(dataset=dataset, batch_size=64, sampler=FixedOrderSampler(indices))
        return loader


    @staticmethod
    def __set_seed():
        seed = 42

        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


    def reload_model(self):
        self.__set_seed()

        self.model = models.resnet34(weights=ResNet34_Weights.DEFAULT if self.pretrained else None).to(self.device)
        self.model.fc = nn.Sequential(
            nn.Linear(self.model.fc.in_features, 1),  # One final output
            nn.Sigmoid()
        ).to(self.device)
        self.__reload_optimizer()


    @staticmethod
    def __measure_metrics(y_true, y_pred):
        res = {}
        for metric_name, metric_fn in [("accuracy", accuracy_score),
                                       ("precision", lambda *args, **kwargs: precision_score(*args, **kwargs, zero_division=0)),  # to suppress warning
                                       ("recall", lambda *args, **kwargs: recall_score(*args, **kwargs, zero_division=0))]:
            res[metric_name] = metric_fn(y_true, y_pred)
        return res


    def fit(self, n_epochs, *, reset=True, show_progress=True):
        if reset:
            self.train_loss_history = []
            self.test_loss_history = []
            self.train_metrics_history = []
            self.reload_model()
        n_batches = len(self.train_loader)

        self.__set_seed()
        #self.model.train()  # Set model to training mode
        with tqdm(range(n_epochs * n_batches), desc="Training", ncols=150, disable=not show_progress) as pbar:
            for epoch in range(1, n_epochs + 1):
                for batch_idx, (batch_images, batch_labels) in enumerate(self.train_loader, 1):
                    batch_images, batch_labels = batch_images.to(self.device), batch_labels.float().reshape(-1, 1).to(self.device)

                    # Forward pass
                    self.model.train()  # Set model to training mode
                    batch_probas = self.model(batch_images)
                    batch_preds = batch_probas >= 0.5
                    batch_loss = self.criterion(batch_probas, batch_labels)

                    # Backward pass
                    self.optimizer.zero_grad()
                    batch_loss.backward()
                    self.optimizer.step()

                    # Record
                    self.train_metrics_history.append(self.__measure_metrics(batch_labels.cpu(), batch_preds.cpu()))
                    self.train_loss_history.append(batch_loss.item())
                    pbar.set_postfix(loss=f"{self.train_loss_history[-1]:.4f}")
                    pbar.update()

                # Evaluate after every epoch
                self.test_loss_history.append(self.predict(show_progress=False)["loss"])


    def plot_train(self, granularity=1):
        train_loss = self.train_loss_history[::granularity]
        x_iter = range(1, len(train_loss) * granularity + 1, granularity)

        n_iters = len(self.train_loss_history)
        n_epochs = len(self.test_loss_history)
        n_batches = int(n_iters / n_epochs) if n_epochs != 0 else 1
        test_loss = self.test_loss_history
        x_epoch = range(n_batches, n_iters + 1, n_batches)

        fig, axes = plt.subplots(1, 2, figsize=(13, 5))
        fig.suptitle(f"Loss and Performance Metrics During Training")

        # Left graph
        ax0 = axes[0]
        ax0.set_title("Loss")
        ax0.plot(x_iter, train_loss, label="train")
        ax0.plot(x_epoch, test_loss, label="test")
        ax0.set_xlabel("iteration")
        ax0.set_ylabel("loss")
        ax0.legend()

        # Right graph
        ax1 = axes[1]
        ax1.set_title("Training Performance")
        if len(self.train_metrics_history) > 0:
            for k in self.train_metrics_history[0].keys():
                ax1.plot(x_iter, [item[k] for item in self.train_metrics_history[::granularity]], label=k)
        ax1.set_xlabel("iteration")
        ax1.set_ylabel("score")
        ax1.set_ylim(0, 1)
        if ax1.get_legend_handles_labels()[0]:
            ax1.legend()

        return fig, axes


    def predict(self, *, show_progress=True):
        self.test_results = {"loss": 0, "probabilities": [], "predictions": []}
        all_labels = []

        self.model.eval()  # Set model to evaluation mode
        with tqdm(range(len(self.test_loader)), desc="Testing", ncols=150, disable=not show_progress) as pbar:
            with torch.no_grad():
                for batch_images, batch_labels in self.test_loader:
                    batch_images, batch_labels = batch_images.to(self.device), batch_labels.float().reshape(-1, 1).to(self.device)

                    # Predict
                    noise = v2.GaussianNoise(sigma=2)
                    batch_images = noise(batch_images)
                    batch_probas = self.model(batch_images)
                    batch_preds = batch_probas >= 0.5

                    # Record
                    all_labels.extend(batch_labels.cpu().numpy())
                    self.test_results["probabilities"].extend(batch_probas.cpu().numpy().ravel().tolist())
                    self.test_results["predictions"].extend(batch_preds.cpu().numpy().astype(np.int8).ravel().tolist())
                    self.test_results["loss"] += self.criterion(batch_probas, batch_labels).item()
                    pbar.update()

        self.test_results["metrics"] = self.__measure_metrics(all_labels, self.test_results["predictions"])
        return self.test_results

# -- Markdown Cell --
# ### Performance without training

# -- Code Cell --
torch.cuda.is_available()

# -- Code Cell --
!pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

# -- Code Cell --
clf = BinaryImageClassifier(lr=0.001)
clf.fit_kfold(n_epochs=10)

# -- Code Cell --
clf = BinaryImageClassifier(lr=0.001)

# -- Code Cell --
# fara noise (0.16597137786448002,
# {'accuracy': 0.98, 'precision': 0.9897959183673469, 'recall': 0.97})


# -- Code Cell --
out = clf.predict()
out["loss"], out["metrics"]

# -- Markdown Cell --
# ## Finetuning

# -- Markdown Cell --
# Finetune this ResNet using data in `data/train`. Use `data/test` as your testing set, and use cross entropy loss. The rest is up to you. Run finetuning that terminates within approx 10 mins. Store the following info every 10 minibatches: loss, precision, recall and accuracy on train and test datasets.

# -- Code Cell --
# Freeze all parameters except those in the fc layer
for param in clf.model.parameters():
    param.requires_grad = False

for param in clf.model.fc.parameters():
    param.requires_grad = True

# -- Markdown Cell --
# > The following cell takes ~3 mins to run on a T4 GPU.

# -- Code Cell --
clf.fit(n_epochs=10)

# -- Code Cell --
clf.plot_train(granularity=10);

# -- Code Cell --
out = clf.predict()
out["loss"], out["metrics"]

# -- Code Cell --
torch.save(clf.model.state_dict(), 'resnet34_finetuned.pth')

# -- Code Cell --
# Training set -> KFold CV -> 5 splituri pentru setul de antreanre 400/100
# finetunati de 5 ori
# augumentare -> augmix

# -- Markdown Cell --
# ## Writeup

# -- Markdown Cell --
# Summarize what you did above, as well as detail the choices you made and why. What was the outcome?

# -- Markdown Cell --
# ---
# 
# A ResNet34 classifies [1000 classes](https://deeplearning.cms.waikato.ac.nz/user-guide/class-maps/IMAGENET/), one of which is chihuahua. Likely because of this, we achieved a perfect score on the test set with minimal fine-tuning.
# 
# We loaded a pretrained ResNet34 model and replaced its final fully connected (FC) layer so that it produces a single output. This output is then passed through a sigmoid activation to clamp the value between 0 and 1.
# 
# We frozen all other layers and trained only the new FC layer. In 10 epochs, both the training and test loss decreased steadily, and finally the model reached 100% accuracy.
# 
# Before moving on, I would like to discuss some mistakes I made while making the code deterministic. I spent a lot of time on repetitive tasks as the hyperparameters obtained did not reproduce the same results, and often the performance was very different!
# 
# 1. I used to call `self.__set_seed()` only in the `__init__` method. However, during the tuning process, I found that the results were not the same each time I called `fit()`. I then realized that I should set the seed every time before training and testing, not just once during initialization.
#    
# 2. When using `DataLoader(dataset, batch_size)`, I was getting batches with all zeros followed by ones. I wanted to achieve a mix of zeros and ones in each batch (which is necessary for the next section). My initial attempt was to use `shuffle=True`, but I later realized that it shuffles the dataset at the start of each epoch, causing different runs to produce different results. Eventually, I came up with a solution by writing a custom sampler and manually shuffling the dataset.
# 
#    ```python
#     class FixedOrderSampler(Sampler):
#     ...
# 
#     class BinaryImageClassifier:
#         ...
#         def load_data(self, ...):
#             ...
#             indices = list(range(len(dataset)))
#             random.shuffle(indices)
#             indices = indices[start:stop:step]
#             
#             loader = DataLoader(dataset=dataset, batch_size=64, sampler=FixedOrderSampler(indices))
#         ...
#     ```
# 
# However, the behavior still differs between Colab and my local Jupyter Notebook for reasons that are unknown to me.

# -- Markdown Cell --
# ## Further analysis

# -- Markdown Cell --
# Pick one aspect about the work done above thus far that you find interesting, investigate it a bit further, and give a short paragraph writeup of what you investigated and how it went.

# -- Markdown Cell --
# ---
# 
# ### Defining the problem
# 
# The original problem was relatively easy to train. To make it more challenging, we impose the following restrictions:
# 
# - Limit the training set to 400 images.
# - Remove pretrained weights (i.e., start ResNet from scratch).
# 
# Our objective is to maximize classification performance (the accuracy, precision, and recall) under these constraints by tuning only three hyperparameters:
# 
# - Learning rate (`lr`)
# - Number of epochs (`n_epochs`)
# - L2 penalty (`l2_penalty`), which corresponds to `weight_decay` in `torch.optim.Adam`
# 
# We do not consider any other modifications beyond these adjustments.

# -- Markdown Cell --
# ### Setup

# -- Code Cell --
clfx = BinaryImageClassifier(lr=0.001, pretrained=False)
clfx.train_loader = clfx.load_data("./train", 0, 400)
clfx.test_loader = clfx.load_data("./train", 400, 500)

len(clfx.train_loader.sampler), len(clfx.test_loader.sampler)

# -- Markdown Cell --
# Here comes the significance of shuffling the images (as mentioned in the "Writeup" section). We don't want the model to train on batches where all images are labeled as zeros.
# 
# Besides, we set aside 100 images from the `./train` folder as the test set. A larger test set would provide more stable and reliable performance estimates.

# -- Markdown Cell --
# ### Tuning

# -- Markdown Cell --
# ### Learning rate

# -- Markdown Cell --
# > The following cell takes ~15 mins to run on a T4 GPU.

# -- Code Cell --
his = []
for lr_exp in range(-5, 0):
    clfx.l2_penalty = 0
    clfx.lr = 10 ** lr_exp
    print(f"Experimenting lr={clfx.lr}...")

    clfx.fit(20)
    his.append((clfx.lr, clfx.train_loss_history, clfx.test_loss_history))

# -- Code Cell --
granularity = 10
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle("Performance on Different Learning Rates")

# Left graph
ax0 = axes[0]
for lr, train_loss, _ in his:
    ax0.plot(range(1, len(train_loss[::granularity]) * granularity + 1, granularity), train_loss[::granularity], label=f"lr={lr}")
ax0.set_xlabel("iteration")
ax0.set_ylabel("loss")
ax0.set_title("Train Loss")
ax0.set_ylim(0, 1)
ax0.legend()

# Right graph
ax1 = axes[1]
for lr, _, test_loss in his:
    ax1.plot(range(1, len(test_loss) + 1), test_loss, label=f"lr={lr}")
ax1.set_xlabel("epoch")
ax1.set_ylabel("loss")
ax1.set_title("Test Loss")
ax1.set_ylim(0, 5)
ax1.legend();

# -- Markdown Cell --
# ### Number of epochs

# -- Markdown Cell --
# Using the best `lr` found from the graph above, we now train for 50 epochs to take a closer look on how the test loss evolves over time.

# -- Markdown Cell --
# > The following cell takes ~7 mins to run on a T4 GPU.

# -- Code Cell --
clfx.lr = 0.0001
clfx.l2_penalty = 0
clfx.fit(50)

# -- Code Cell --
fig, axes = clfx.plot_train(10)

min_idx = np.argmin(clfx.test_loss_history)
min_x = (min_idx + 1) * len(clfx.train_loader)
min_y = clfx.test_loss_history[min_idx]

axes[0].scatter(min_x, min_y, color='red')
axes[0].annotate(f'min loss: {min_y:.4f}\n(n_epochs={min_idx + 1})', (min_x, min_y), textcoords="offset points", xytext=(5,7));

# -- Markdown Cell --
# ### L2 penalty

# -- Markdown Cell --
# After some empirical testing, it appears that adding an L2 penalty did not improve the scores. Therefore, we set it to zero. Our final result is:

# -- Code Cell --
clfx.lr = 0.0001
clfx.l2_penalty = 0
clfx.fit(16)

out = clfx.predict()
out["loss"], out["metrics"]

# -- Code Cell --
torch.save(clfx.model.state_dict(), 'resnet34_manual400.pth')

# -- Markdown Cell --
# ## Auto tuning
# 
# In this section, we utilize Optuna to automatically tune our hyperparameters. Optuna automates hyperparameter tuning using Bayesian optimization. It smartly selects hyperparameter configurations based on past trial outcomes, focusing on promising regions of the search space.

# -- Code Cell --
!pip install optuna

# -- Code Cell --
import optuna
import optuna.visualization as vis

# -- Code Cell --
def objective(trial):
    lr = trial.suggest_float('lr', 1e-5, 1e-1, log=True)
    n_epochs = trial.suggest_int('n_epochs', 1, 20)
    l2_penalty = trial.suggest_float('l2_penalty', 1e-6, 1e-1, log=True)

    clfy = BinaryImageClassifier(lr=lr, l2_penalty=l2_penalty, pretrained=False)

    clfy.train_loader = clfy.load_data("./train", 0, 400)
    clfy.test_loader = clfy.load_data("./train", 400, 500)

    clfy.fit(n_epochs)
    out = clfy.predict()
    return *out["metrics"].values(), out["loss"]

# -- Markdown Cell --
# > The following cell takes ~1 hour to run on a T4 GPU.

# -- Code Cell --
study = optuna.create_study(directions=["maximize", "maximize", "maximize", "minimize"],
                            sampler=optuna.samplers.TPESampler(seed=42))
study.optimize(objective, n_trials=30)

# -- Code Cell --
vis.plot_optimization_history(
    study,
    target=lambda t: t.values[0],
    target_name="Accuracy"
)

# -- Code Cell --
vis.plot_slice(
    study,
    target=lambda t: t.values[0],
    target_name="Accuracy"
)

# -- Code Cell --
vis.plot_param_importances(
    study,
    target=lambda t: t.values[0],
    target_name="Accuracy"
)

# -- Code Cell --
study.best_trials

# -- Code Cell --
best_params = study.trials[19].params
best_params

# -- Code Cell --
study.trials[19].values

# -- Code Cell --
clfy = BinaryImageClassifier(lr=best_params["lr"], l2_penalty=best_params["l2_penalty"], pretrained=False)

clfy.train_loader = clfy.load_data("./train", 0, 400)
clfy.test_loader = clfy.load_data("./train", 400, 500)

clfy.fit(best_params["n_epochs"])

# -- Code Cell --
clfy.plot_train(10);

# -- Code Cell --
out = clfy.predict()
out["loss"], out["metrics"]

# -- Code Cell --
torch.save(clfy.model.state_dict(), 'resnet34_auto400.pth')

# -- Markdown Cell --
# ## Conclusion
# 
# | | Parameters | Accuracy | Precision | Recall | Loss |
# | --- | --- | --- | --- | --- | --- |
# | Manual | `lr=1.000e-4, n_epochs=16, l2_penalty=0.000` | 0.9300 | 0.9375 | 0.9184 | 0.4191 |
# | Optuna | `lr=2.351e-4, n_epochs=20, l2_penalty=1.917e-4` | 0.9700 | 0.9792 | 0.9592 | 0.3445 |
# 
# (Correct to 4 sig. fig.)
# 
# The results demonstrate that the Optuna-optimized model outperformed the manually tuned counterpart across all evaluation metrics, achieving higher accuracy, precision, and recall, while also reducing the loss.
# 
# This highlights the effectiveness of automated hyperparameter optimization in improving model performance and reducing the need for extensive manual experimentation. Also, the manual approach of "tuning one parameter at a time" may be suboptimal, as hyperparameters often interact in non-orthogonal ways: changes in one can influence the optimal values of others.