from typing import Callable, Dict, List, Tuple

import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader, TensorDataset

from utils import load_dataset, problem

LossFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]


@problem.tag("hw3-A")
def J_loss(logits: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Compute the joint negative log-likelihood loss.

    This is the sum of per-example negative log-likelihoods for the
    multinomial logistic regression model.

    Args:
        logits: FloatTensor of shape (n, k). Raw class scores.
        y: LongTensor of shape (n,). Class labels.

    Returns:
        A scalar tensor containing the loss.
    """
    raise NotImplementedError("Your Code Goes Here")


@problem.tag("hw3-A")
def L_loss(logits: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Compute the standard average cross-entropy loss.

    This is the mean per-example negative log-likelihood. You may use
    torch.nn.functional.cross_entropy or write the softmax/log expression
    directly.

    Args:
        logits: FloatTensor of shape (n, k). Raw class scores.
        y: LongTensor of shape (n,). Class labels.

    Returns:
        A scalar tensor containing the loss.
    """
    raise NotImplementedError("Your Code Goes Here")


@problem.tag("hw3-A")
def accuracy(logits: torch.Tensor, y: torch.Tensor) -> float:
    """
    Compute classification accuracy.

    Args:
        logits: FloatTensor of shape (n, k). Raw class scores.
        y: LongTensor of shape (n,). Class labels.

    Returns:
        Fraction of examples classified correctly.
    """
    raise NotImplementedError("Your Code Goes Here")


@problem.tag("hw3-A")
def train(
    x_train: torch.Tensor,
    y_train: torch.Tensor,
    loss_fn: LossFunction,
    *,
    learning_rate: float = 0.1,
    epochs: int = 50,
    batch_size: int = 256,
) -> Tuple[torch.Tensor, List[float]]:
    """
    Train a multinomial logistic regression model on MNIST.

    Args:
        x_train: FloatTensor of shape (n, d).
        y_train: LongTensor of shape (n,).
        loss_fn: Either J_loss or L_loss.
        learning_rate: Step size for gradient descent.
        epochs: Number of training epochs.
        batch_size: Mini-batch size.

    Returns:
        A tuple (W, losses), where W has shape (k, d) and losses stores one
        average training loss per epoch.
    """
    dataset = TensorDataset(x_train, y_train)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    _ = loader
    raise NotImplementedError("Your Code Goes Here")


@problem.tag("hw3-A", start_line=5)
def main() -> Dict[str, Dict[str, float]]:
    """
    Train multinomial logistic regression models with J_loss and L_loss.

    For each loss, this function should:
        1. Load MNIST.
        2. Train a multinomial logistic regression model.
        3. Plot training loss vs. epoch.
        4. Report training and test accuracy.

    Returns:
        A dictionary mapping loss names to accuracy summaries.
    """
    (x_train, y_train), (x_test, y_test) = load_dataset("mnist")
    x_train = torch.from_numpy(x_train).float()
    y_train = torch.from_numpy(y_train).long()
    x_test = torch.from_numpy(x_test).float()
    y_test = torch.from_numpy(y_test).long()

    raise NotImplementedError("Your Code Goes Here")


if __name__ == "__main__":
    main()
