"""
Utilities to be used along with the deep model
"""
from typing import Union

import torch
from torch import nn
from proj5_code.simple_net import SimpleNet
from proj5_code.simple_net_final import SimpleNetFinal
from proj5_code.my_resnet import MyResNet18


def compute_accuracy(logits: torch.Tensor, labels: torch.Tensor) -> float:
    """
    Compute the accuracy given the prediction logits and the ground-truth
    labels.

    Args:
        logits: the (batch_size, num_classes) output of the forward pass
            through the model. For K classes, logits[k] (where 0 <= k < K)
            corresponds to the log-odds of class `k` being the correct one.
        labels: the ground truth label of shape (batch_size) for each instance
            in the batch

    Returns:
        accuracy: The accuracy of the predicted logits
            (number of correct predictions / total number of examples)
    """
    batch_accuracy = 0.0
    ###########################################################################
    # Student code begins
    ###########################################################################

    predicted_classes = torch.argmax(logits, dim=1)
    correct_labels = predicted_classes == labels

    batch_accuracy = float(torch.sum(correct_labels)) / len(labels)

    ###########################################################################
    # Student code ends
    ###########################################################################
    return batch_accuracy


def compute_loss(
    model: Union[SimpleNet, SimpleNetFinal, MyResNet18],
    model_output: torch.Tensor,
    target_labels: torch.Tensor,
    is_normalize: bool = True,
) -> torch.Tensor:
    """
    Computes the loss between the model output and the target labels.

    Args:
        model: a model (which inherits from nn.Module)
        model_output: the raw scores output by the net
        target_labels: the ground truth class labels
        is_normalize: bool flag indicating that loss should be divided by the
            batch size
    Returns:
        the loss value
    """
    loss = None
    ############################################################################
    # Student code begin
    ############################################################################

    # target_labels = target_labels.unsqueeze(dim=1)

    # loss = model.loss_criterion(model_output, target_labels.float())

    (batch_size, n_label) = model_output.shape

    x_labels = torch.arange(0, batch_size)

    target = torch.zeros(model_output.shape)
    target[x_labels, target_labels] = 1

    loss = model.loss_criterion(model_output, target.float().cuda())

    if is_normalize:
        (batch_size, output_size) = model_output.shape
        loss = loss / batch_size


    loss = loss.to(torch.float32)

    ############################################################################
    # Student code end
    ############################################################################
    return loss


def save_trained_model_weights(
    model: Union[SimpleNet, SimpleNetFinal, MyResNet18, nn.Module], out_dir: str
) -> None:
    """Saves the weights of a trained model along with class name

    Args:
        model: The model to be saved
        out_dir: The path to the folder to store the save file in
    """
    class_name = model.__class__.__name__
    state_dict = model.state_dict()

    assert class_name in set(
        ["SimpleNet", "SimpleNetFinal", "MyResNet18"]
    ), "Please save only supported models"

    save_dict = {"class_name": class_name, "state_dict": state_dict}
    torch.save(save_dict, f"{out_dir}/trained_{class_name}_final.pt")
