r"""This module defines a criterion wrapper to make PyTorch criteria
compatible with ``gravitorch.models.VanillaModel``."""

from __future__ import annotations

__all__ = ["VanillaLoss"]


from torch import Tensor
from torch.nn import Module

from gravitorch import constants as ct
from gravitorch.nn import setup_module


class VanillaLoss(Module):
    r"""Implements a wrapper to make compatible most of the PyTorch loss
    functions (a.k.a. criterion) with
    ``gravitorch.models.VanillaModel``.

    This wrapper assumes the loss function has two inputs:

        - a tensor of prediction.
        - a tensor of target

    The shape and type of the tensors depend on the loss function used.

    Args:
    ----
        criterion (``torch.nn.Module`` or dict): Specifies the loss
            function (a.k.a. criterion) or its configuration.
        prediction_key (str): Specifies the prediction key.
            Default: ``"prediction"``.
        target_key (str): Specifies the target key.
            Default: ``"target"``.

    Example usage:

    .. code-block:: pycon

        >>> import torch
        >>> from gravitorch.models.criteria import VanillaLoss
        >>> # Initialization with a nn.Module
        >>> criterion = VanillaLoss(criterion=torch.nn.MSELoss())
        >>> criterion
        VanillaLoss(
          (criterion): MSELoss()
        )
        >>> # Initialization with a config
        >>> criterion = VanillaLoss(criterion={"_target_": "torch.nn.MSELoss"})
        >>> criterion
        VanillaLoss(
          (criterion): MSELoss()
        )
        >>> # Customize keys.
        >>> criterion = VanillaLoss(
        ...     criterion=torch.nn.MSELoss(),
        ...     prediction_key="my_prediction",
        ...     target_key="my_target",
        ... )
        >>> criterion
        VanillaLoss(
          (criterion): MSELoss()
        )
        >>> loss = criterion(
        ...     net_out={"my_prediction": torch.rand(2, 4)}, batch={"my_target": torch.rand(2, 4)}
        ... )
        >>> loss  # doctest: +ELLIPSIS
        {'loss': tensor(...)}
    """

    def __init__(
        self,
        criterion: Module | dict,
        prediction_key: str = ct.PREDICTION,
        target_key: str = ct.TARGET,
    ) -> None:
        super().__init__()
        self.criterion = setup_module(criterion)
        self._prediction_key = prediction_key
        self._target_key = target_key

    def forward(self, net_out: dict, batch: dict) -> dict[str, Tensor]:
        r"""Computes the loss value given the network output and the
        batch.

        Args:
        ----
            net_out (dict): Specifies the network output
                which contains the prediction.
            batch (dict): Specifies the batch which contains the target.

        Returns:
        -------
            dict: a dict with the loss value.
        """
        prediction = self._get_prediction_from_net_out(net_out)
        target = self._get_target_from_batch(batch)
        return {ct.LOSS: self.criterion(prediction, target)}

    def _get_prediction_from_net_out(self, net_out: dict) -> Tensor:
        r"""Gets the prediction from the network output.

        Args:
        ----
            net_out (dict): Specifies the network output
                which contains the prediction.

        Returns:
        -------
            ``torch.Tensor``: The prediction
        """
        return net_out[self._prediction_key]

    def _get_target_from_batch(self, batch: dict) -> Tensor:
        r"""Gets the target from the batch. The target is the tensor with
        the key 'target'.

        Args:
        ----
            batch (dict): Specifies the batch which contains the target.

        Returns:
        -------
            ``torch.Tensor``: The target
        """
        return batch[self._target_key]
