Non-Normal Distribution

######################### DO NOT CHANGE THIS CELL ##########################

def add_normal_noise(image, mean=0, std=0.2):
    """Adds normal (Gaussian) noise to the image."""
    noise = torch.distributions.normal.Normal(mean, std)
    noise = noise.sample(image.size())
    noisy_image = image + noise
    return torch.clamp(noisy_image, 0.0, 1.0)

######################### DO NOT CHANGE THIS CELL ##########################

def add_uniform_noise(image, low=-0.5, high=0.5):
    """Adds uniform noise to the image."""
    noise = torch.empty(image.size()).uniform_(low, high)
    noisy_image = image + noise
    return torch.clamp(noisy_image, 0.0, 1.0)

######################### DO NOT CHANGE THIS CELL ##########################

FINAL_EVALUATION_MODE = False  # We will set this flag to True during grading.

######################### DO NOT CHANGE THIS CELL ##########################

import os
import pickle
import random
from typing import Dict, List, Tuple
from collections.abc import Callable
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset


DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

assert torch.cuda.is_available(), "CUDA not available!"

print("Device:", DEVICE)

######################### DO NOT CHANGE THIS CELL ##########################

seed = 42

os.environ["PYTHONHASHSEED"] = str(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

######################### DO NOT CHANGE THIS CELL ##########################
# Cell containing helper functions for visualising the results

def plot_samples(dataset: Dataset, num_images: int = 6, title: str = "") -> None:
    """
    Function that displays examples of original and noisy images

    Arguments:
        dataset (Dataset): The dataset.
        num_images (int): Number of images to display.
    """
    fig, axs = plt.subplots(2, num_images, figsize=(2 * num_images, 4))
    fig.suptitle(title, fontsize=10)

    for i in range(num_images):
        sample = dataset[i]
        original_image = sample["original"]
        noised_image = sample["noised"]
        label = sample["label"]
        params = sample.get("params", None)

        column_title = f"Example {i+1}\nLabel: {label.item():.0f}"

        if params is not None:
            if label.item() == 0:
                column_title += f"\nμ: {params[0].item():.2f}\nσ: {params[1].item():.2f}"
            else:
                column_title += f"\nlow: {params[0].item():.2f}\nhigh: {params[1].item():.2f}"
        else:
            if label.item() == 0:
                column_title += "\nμ: None\nσ: None"
            else:
                column_title += "\nlow: None\nhigh: None"

        column_title += "\n\nOriginal"

        axs[0, i].set_title(column_title, fontsize=8, pad=5)
        axs[0, i].imshow(original_image.squeeze(), cmap="gray")
        axs[0, i].axis("off")

        axs[1, i].set_title("Noisy", fontsize=8, pad=5)
        axs[1, i].imshow(noised_image.squeeze(), cmap="gray")
        axs[1, i].axis("off")

    plt.tight_layout()
    plt.show()


def plot_results(model: nn.Module, examples: Dict, num_images: int = 6) -> None:
    """
    Function that displays examples of a noisy image and the same image denoised by the model

    Arguments:
        model (nn.Module): Image denoising model.
        examples (dict): Dictionary containing example images.
        num_images (int): Number of images to display.
    """
    model.eval()
    noisy_images = examples["noised"][:num_images].to(DEVICE)
    clean_images = examples["original"][:num_images]
    label = examples["label"][:num_images]
    params = examples["params"][:num_images]
    mean_real = params[:, 0].view(-1, 1)
    std_real = params[:, 1].view(-1, 1)

    with torch.no_grad():
        output_images, predictions, mean_pred, std_pred = model(noisy_images.to(DEVICE))

    fig, axs = plt.subplots(3, num_images, figsize=(2 * num_images, 6))

    for i in range(num_images):
        column_title = (
            f"Example {i+1}\n"
            f"Label: {float(predictions[i].item() > 0.5):.0f}/{label[i].item():.0f}\n"
        )

        if label[i].item() == 0:
            column_title += (
                f"μ: {mean_pred[i].item():.2f}/{mean_real[i].item():.2f}\n"
                f"σ: {std_pred[i].item():.2f}/{std_real[i].item():.2f}\n"
            )
        else:
            column_title += "\n\n"

        column_title += "\nNoisy"

        axs[0, i].set_title(column_title, fontsize=8, pad=5)
        axs[0, i].imshow(noisy_images[i].cpu().squeeze(), cmap="gray")
        axs[0, i].axis("off")

        axs[1, i].set_title("Denoised", fontsize=8, pad=5)
        axs[1, i].imshow(output_images[i].cpu().squeeze(), cmap="gray")
        axs[1, i].axis("off")

        axs[2, i].set_title("Original", fontsize=8, pad=5)
        axs[2, i].imshow(clean_images[i].squeeze(), cmap="gray")
        axs[2, i].axis("off")

    fig.text(0.5, 0.01, "Format: Prediction/Label", ha='center', fontsize=10)

    plt.tight_layout()
    plt.subplots_adjust(top=0.85, bottom=0.1)
    plt.show()

######################### DO NOT CHANGE THIS CELL ##########################
# Cell containing helper functions for preparing the data.

class NoisedDataset(Dataset):
    """
    Dataset loaded from a pickle file.

    Arguments:
        file_path (str): Path to the pickle file containing the data.
        transform (callable, optional): Transformations applied to the images and labels.
    """
    def __init__(self, pickle_file, transform=None):
        self.transform = transform

        with open(pickle_file, 'rb') as f:
            self.data = pickle.load(f)

        self.has_params = 'params' in self.data[0]

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:
        sample = self.data[idx]
        original_image = sample['original']
        noised_image = sample['noised']
        label = sample['label']

        if self.has_params:
            data = {'params': sample['params']}
        else:
            data = {}

        if self.transform:
            original_image = self.transform(original_image)
            noised_image = self.transform(noised_image)

        data.update({
            'original': original_image,
            'noised': noised_image,
            'label': label
        })

        return data


def setup_data(
        train_transform: Callable | None = None,
        val_transform: Callable | None = None,
        root: str = './'
    ) -> Tuple[Dataset, Dataset]:
    """
    Prepares the training and validation datasets, downloading them if necessary.

    Arguments:
        train_transform (callable, optional): Augmentations for the training set.
        val_transform (callable, optional): Augmentations for the validation set.
        root (str, optional): Base directory for the data files.

    Returns:
        tuple: Datasets (train_ds, val_ds).
    """
    if train_transform is None:
        train_transform = transforms.Compose([transforms.ToTensor()])
    if val_transform is None:
        val_transform = transforms.Compose([transforms.ToTensor()])

    train_file = root+'train.pkl'
    val_file = root+'val.pkl'

    if not os.path.exists(root):
        os.makedirs(root)

    train_ds = NoisedDataset(train_file, transform=train_transform)
    val_ds = NoisedDataset(val_file, transform=val_transform)

    return train_ds, val_ds

######################### DO NOT CHANGE THIS CELL ##########################

train_ds, val_ds = setup_data(root="./")

if not FINAL_EVALUATION_MODE:
    print("Number of images in the training set:", len(train_ds), ", number of images in the validation set:", len(val_ds))
    plot_samples(train_ds, num_images=6, title="Training set")
    plot_samples(val_ds, num_images=6, title="Validation set")

######################### DO NOT CHANGE THIS CELL ##########################
# Cell containing helper functions for computing the model's metric values

def compute_psnr(input_image: torch.Tensor, target_image: torch.Tensor) -> torch.Tensor:
    """
    Function that computes the PSNR between two images.

    Arguments:
        input_image (torch.Tensor): The first image.
        target_image (torch.Tensor): The second image.

    Returns:
        torch.Tensor: The PSNR value.
    """
    mse = F.mse_loss(input_image, target_image)
    if mse == 0:
        return 100
    return 10 * torch.log10(1 / mse)

def model_eval(model: nn.Module, dataloader: DataLoader, device: str = DEVICE) -> Tuple[float, float, float, float]:
    """"
    Function for evaluating the model on a dataset.

    Arguments:
        model (nn.Module): The model to evaluate.
        dataloader (DataLoader): DataLoader with the evaluation data.
        device (str, optional): The device on which the evaluation is to be run.

    Returns:
        tuple: A tuple containing the metric values (PSNR, accuracy, MSE for parameter 1, MSE for parameter 2).
    """
    model.eval()
    model.to(device)

    # Initialise the variables that store the results
    psnr = 0
    correct = 0
    mean_mse = 0
    std_mse = 0

    total_samples = 0
    total_label0_samples = 0

    with torch.no_grad():
        for data in dataloader:
            noised_images = data["noised"].to(device)
            original_images = data["original"].to(device)
            labels = data["label"].to(device)
            params = data["params"].to(device)
            batch_size = len(labels)

            mean_real = params[:, 0].view(-1, 1)
            std_real = params[:, 1].view(-1, 1)

            output_images, labels_pred, mean_pred, std_pred = model(noised_images)

            # Compute the classification accuracy
            correct += ((labels_pred >= 0.5).float().view(-1) == labels).sum().item()

            # Compute the PSNR
            psnr += compute_psnr(output_images, original_images) * batch_size

            # Compute the MSE for the parameters when the label is 0
            label0_mask = (labels == 0)
            num_label0 = label0_mask.sum().item()

            if num_label0 > 0:
                mean_mse += F.mse_loss(mean_pred[label0_mask], mean_real[label0_mask], reduction='sum')
                std_mse += F.mse_loss(std_pred[label0_mask], std_real[label0_mask], reduction='sum')

            total_samples += batch_size
            total_label0_samples += num_label0

    # Compute the mean values of the metrics
    psnr /= total_samples
    accuracy = correct / total_samples
    mean_mse /= total_label0_samples
    std_mse /= total_label0_samples

    return psnr.item(), accuracy, mean_mse.item(), std_mse.item()

######################### DO NOT CHANGE THIS CELL ##########################
# Cell containing helper functions for scoring your solution

def calculate_score(
    psnr: float, accuracy: float, mean_mse: float, std_mse: float
) -> Tuple[float, float, float, float]:
    """
    Function that computes the points for the task from the model's metrics.

    Arguments:
        psnr (float): The PSNR value.
        accuracy (float): The classification accuracy.
        mean_mse (float): MSE for parameter 1.
        std_mse (float): MSE for parameter 2.

    Returns:
        tuple: A tuple containing the points for the task (PSNR, accuracy, MSE for parameter 1, MSE for parameter 2).
    """

    def scale(x, lower=0.0, upper=1.0, max_points=1.0):
        scaled = min(max(x, lower), upper)
        return (scaled - lower) / (upper - lower) * max_points

    accuracy_score = scale(accuracy, lower=0.5, upper=0.95)
    psnr_score = scale(psnr, lower=10.0, upper=16.0)

    mean_score = 0.0
    if mean_mse < 0.005:
        mean_score = 1.0

    std_score = 0.0
    if std_mse < 0.005:
        std_score = 1.0

    return psnr_score, accuracy_score, mean_score, std_score


def grade_solution(model: nn.Module, dataloader: DataLoader) -> float:
    """
    Function that scores the model on the validation set.

    Arguments:
        model (nn.Module): The model to score.
        dataloader (DataLoader): DataLoader with the scoring data.

    Returns:
        float: The number of points for the task.
    """
    psnr, accuracy, mean_mse, std_mse = model_eval(model, dataloader)
    psnr_score, accuracy_score, mean_score, std_score = calculate_score(
        psnr, accuracy, mean_mse, std_mse
    )
    score = round(
        psnr_score * 25 + accuracy_score * 25 + mean_score * 25 + std_score * 25
    )

    # Round to an integer, range [0, 100]
    score = round(score)

    print(
        f"Metrics on the validation set\n"
        f"psnr: {psnr:.2f}, accuracy: {accuracy:.2f}, mean_mse: {mean_mse:.6f}, std_mse: {std_mse:.6f}\n"
    )
    print(
        f"Partial points for the task\n"
        f"psnr: {(psnr_score * 25):.2f}, accuracy: {(accuracy_score * 25):.2f}, mean_mse: {(mean_score * 25):.2f}, std_mse: {(std_score * 25):.2f}\n"
    )
    print(f"Estimated number of points for the task: {score}")
    return score

# definitions of the augmentations for the training and validation sets. The default None means no augmentation
train_transform = None
val_transform = None

# batch size
BATCH_SIZE: int = 64

train_ds, val_ds = setup_data(train_transform, val_transform, root="./")

# create and train your model here

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        """
        The transformation performed by the model.

        Arguments:
            x (torch.Tensor): Input image [B, 1, H, W].

        Returns:
            The result, returned as a tuple:
            torch.Tensor: Output image [B, 1, H, W].
            torch.Tensor: Classification predictions [B, 1].
            torch.Tensor: Parameter 1 [B, 1].
            torch.Tensor: Parameter 2 [B, 1].
        """
        device = x.device
        return (
            torch.rand_like(x, device=device),
            torch.rand(x.shape[0], 1, device=device),
            torch.randn(x.shape[0], 1, device=device),
            torch.randn(x.shape[0], 1, device=device)
        )


def train_model() -> Model:
    """Create and train the model"""
    return Model().to(DEVICE)


your_model = train_model()

######################### DO NOT CHANGE THIS CELL ##########################

val_dataloader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False)

if not FINAL_EVALUATION_MODE:
    grade_solution(your_model, val_dataloader)
    examples = next(iter(val_dataloader))
    plot_results(your_model, examples, num_images=6)

######################### DO NOT CHANGE THIS CELL ##########################

if FINAL_EVALUATION_MODE:
    import cloudpickle

    # If the model has parameters, put it in evaluation mode and move it to the CPU
    if list(your_model.parameters()):
        your_model.eval()
        your_model.cpu()

    OUTPUT_PATH = "file_output"
    FUNCTION_FILENAME = "your_model.pkl"
    FUNCTION_OUTPUT_PATH = os.path.join(OUTPUT_PATH, FUNCTION_FILENAME)

    if not os.path.exists(OUTPUT_PATH):
        os.makedirs(OUTPUT_PATH)

    with open(FUNCTION_OUTPUT_PATH, "wb") as f:
        cloudpickle.dump(your_model, f)