"""Module implements MPT Model for Sequence Classification."""
import logging
from collections.abc import MutableMapping
from typing import TYPE_CHECKING, cast

import torch
from composer.models import HuggingFaceModel
from llmfoundry import ComposerMPTCausalLM
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from torchmetrics import Metric
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,
    SequenceClassifierOutputWithPast,
)

if TYPE_CHECKING:
    from llmfoundry.models import MPTForCausalLM
    from transformers.configuration_utils import PretrainedConfig

logger = logging.getLogger(__name__)

TWO_D_TENSOR_DIMENSION = 2
ONE_D_TENSOR_DIMENSION = 1


class MPTForSequenceClassification(HuggingFaceModel):
    """MPT Model for Sequence Classification.

    This class implements a sequence classification model using the MPT (Multi-Document
    Processing Transformer) architecture. It extends the HuggingFaceModel class and
    provides functionality for training and evaluating sequence classification tasks.

    Attributes
    ----------
    num_labels : int
        The number of labels for classification.
    problem_type : str | None
        The type of classification problem (e.g., "regression",
        "single_label_classification", "multi_label_classification").
    pad_token_id : int
        The ID of the padding token.
    transformer : ComposerMPTCausalLM
        The transformer model used for sequence classification.
    score : nn.Linear
        A linear layer to compute the classification scores.

    Methods
    -------
    __init__(
        model,
        num_labels,
        hidden_size,
        pad_token_id,
        problem_type=None,
        train_metrics=None,
        eval_metrics=None
    )
        Initialize the MPTForSequenceClassification model.
    forward(batch) -> tuple[torch.Tensor] | SequenceClassifierOutputWithPast
        Perform a forward pass on the input batch and return the classification output.
    eval_forward(batch, outputs=None)
        Perform a forward pass on the input batch during evaluation and return the
        logits.

    Parameters
    ----------
    model : ComposerMPTCausalLM
        The transformer model to be used for sequence classification.
    num_labels : int
        The number of labels for classification.
    hidden_size : int
        The size of the hidden layer in the transformer model.
    pad_token_id : int
        The ID of the padding token.
    problem_type : str | None, optional
        The type of classification problem (default is None).
    train_metrics : list[Metric] | None, optional
        A list of metrics to be used during training (default is None).
    eval_metrics : list[Metric] | None, optional
        A list of metrics to be used during evaluation (default is None).

    Example
    -------
    >>> model = ComposerMPTCausalLM(...)
    >>> classifier = MPTForSequenceClassification(
    ...     model=model,
    ...     num_labels=2,
    ...     hidden_size=768,
    ...     pad_token_id=0,
    ...     problem_type="single_label_classification",
    ...     train_metrics=[...],
    ...     eval_metrics=[...]
    ... )
    >>> batch = {"input_ids": ..., "labels": ...}
    >>> output = classifier.forward(batch)
    >>> print(output)

    """

    def __init__(  # noqa: PLR0913, PLR0917
        self,
        model: ComposerMPTCausalLM,
        num_labels: int,
        hidden_size: int,
        pad_token_id: int | None,
        problem_type: str | None = None,
        train_metrics: list[Metric] | None = None,
        eval_metrics: list[Metric] | None = None,
    ) -> None:
        """Initialize the MPTForSequenceClassification model.

        Parameters
        ----------
        model : ComposerMPTCausalLM
            The transformer model to be used for sequence classification.
        num_labels : int
            The number of labels for classification.
        hidden_size : int
            The size of the hidden layer in the transformer model.
        pad_token_id : int | None
            The ID of the padding token.
        problem_type : str | None, optional
            The type of classification problem (default is None).
        train_metrics : list[Metric] | None, optional
            A list of metrics to be used during training (default is None).
        eval_metrics : list[Metric] | None, optional
            A list of metrics to be used during evaluation (default is None).

        """
        self.num_labels = num_labels
        self.problem_type = problem_type
        self.pad_token_id = pad_token_id

        super().__init__(
            model=model.model,
            tokenizer=model.tokenizer,
            use_logits=True,
            metrics=train_metrics,
            eval_metrics=eval_metrics,
            shift_labels=cast("MPTForCausalLM", model.model).transformer.shift_labels,
            allow_embedding_resizing=True,
        )
        self.transformer = model
        self.score = nn.Linear(hidden_size, num_labels, bias=False)

    def forward(  # noqa: C901, PLR0912, PLR0914
        self,
        batch: MutableMapping,
    ) -> tuple[torch.Tensor, ...] | SequenceClassifierOutputWithPast:
        """Perform a forward pass on the input batch and return a classification output.

        This function processes a batch of input data, performs a forward pass using the
        transformer model, and returns the classification output. It supports different
        problem types, including regression, single-label classification, and multi-
        label classification. The function also handles padding tokens and computes the
        loss if labels are provided.

        Parameters
        ----------
        batch : MutableMapping
            A dictionary containing the input data for the model. It must include the
            following keys:
            - "input_ids" (torch.Tensor): The input token IDs.
            - "inputs_embeds" (torch.Tensor): The input embeddings.
            - "labels" (torch.Tensor): The labels for the input data.
            - "return_dict" (bool): Whether to return a dictionary of outputs.

        Returns
        -------
        tuple[torch.Tensor, ...] | SequenceClassifierOutputWithPast
            The classification output. If `return_dict` is True, a
            SequenceClassifierOutputWithPast object is returned. Otherwise, a tuple
            containing the pooled logits and additional outputs is returned. If labels
            are provided, the loss is included in the output.

        Example
        -------
        >>> batch = {
        ...     "input_ids": torch.tensor([[101, 102, 103], [104, 105, 106]]),
        ...     "inputs_embeds": None,
        ...     "labels": torch.tensor([1, 0]),
        ...     "return_dict": True
        ... }
        >>> model = MPTForSequenceClassification(...)
        >>> output = model.forward(batch)
        >>> print(output)

        Raises
        ------
        ValueError
            If the batch size is greater than 1 and no padding token is defined.


        """
        input_ids = batch.get("input_ids", None)
        inputs_embeds = batch.get("inputs_embeds", None)
        labels = batch.pop("labels", None)
        return_dict = batch.get("return_dict", None)
        output_hidden_states = cast(
            "PretrainedConfig",
            self.transformer.model.config,
        ).output_hidden_states

        return_dict = (
            return_dict
            if return_dict is not None
            else self.transformer.config.use_return_dict
        )
        transformer_outputs: CausalLMOutputWithPast = self.transformer.model(
            **batch,
            output_hidden_states=output_hidden_states,
        )
        assert transformer_outputs is not None
        assert transformer_outputs.hidden_states is not None

        hidden_states, logits = transformer_outputs.hidden_states[-1]
        logits = cast("torch.Tensor", self.score(hidden_states))

        batch_size: int | None = None
        if input_ids is not None:
            batch_size = input_ids.shape[0]
        elif inputs_embeds is not None:
            batch_size = inputs_embeds.shape[0]
        assert batch_size is not None

        if self.pad_token_id is None and batch_size != 1:
            msg = "Cannot handle batch sizes > 1 if no padding token is defined."
            raise ValueError(
                msg,
            )
        if self.pad_token_id is None:
            sequence_lengths = -1
        elif input_ids is not None:
            # If no pad token has been found, use modulo instead of reverse indexing for
            # ONNX compatibility
            sequence_lengths_tensor = (
                torch.eq(input_ids, self.pad_token_id).int().argmax(-1) - 1
            )
            sequence_lengths_tensor %= input_ids.shape[-1]
            sequence_lengths = int(sequence_lengths_tensor.to(logits.device).item())
        else:
            sequence_lengths = -1

        pooled_logits, loss = (
            logits[
                torch.arange(batch_size, device=logits.device),
                sequence_lengths,
            ],
            None,
        )

        if labels is not None:
            if self.problem_type is None:
                if self.num_labels == 1:
                    self.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype in {torch.long, torch.int}):
                    self.problem_type = "single_label_classification"
                else:
                    self.problem_type = "multi_label_classification"

            if self.problem_type == "regression":
                loss_fct = MSELoss()

                loss = (
                    loss_fct(pooled_logits.squeeze(), labels.squeeze())
                    if self.num_labels == 1
                    else loss_fct(pooled_logits, labels)
                )

            elif self.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(pooled_logits, labels)
            elif self.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels)
        if not return_dict:
            output = (
                pooled_logits,
                *cast("tuple[torch.Tensor, ...]", transformer_outputs[1:]),
            )
            return ((loss, *output)) if loss is not None else output

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=cast("torch.FloatTensor", pooled_logits),
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

    @staticmethod
    def _interpret_output(
        output: SequenceClassifierOutputWithPast | tuple[torch.Tensor, ...],
    ) -> torch.Tensor:
        """Interpret the model output and return the logits.

        This function checks the type of the output and extracts the logits from it.
        If the output is of type `SequenceClassifierOutputWithPast`, it retrieves the
        logits directly. If the output is a tuple, it checks the shape of the first
        element to determine whether to return the first element or the second element
        as the logits. It also handles the case where the logits are a 2D tensor with
        a single class, in which case it squeezes the second dimension.

        Parameters
        ----------
        output : SequenceClassifierOutputWithPast | tuple[torch.Tensor, ...]
            The output from the model, which can either be a
            `SequenceClassifierOutputWithPast` object or a tuple containing the logits.

        Returns
        -------
        torch.Tensor
            The logits extracted from the output. If the output is a
            `SequenceClassifierOutputWithPast`, the logits are returned directly. If it
            is a tuple, the appropriate tensor is selected based on its shape.

        """
        if isinstance(output, SequenceClassifierOutputWithPast):
            ret_output = output.logits
            assert ret_output is not None
        else:
            assert isinstance(output, tuple)
            assert len(output) > 1
            ret_output = output[1] if len(output[0].shape) == 0 else output[0]

        # If we are in the single class case, then remove the classes dimension
        if (
            ret_output.ndim == TWO_D_TENSOR_DIMENSION
            and ret_output.shape[1] == ONE_D_TENSOR_DIMENSION
        ):
            ret_output = ret_output.squeeze(dim=1)

        return ret_output

    def eval_forward(
        self,
        batch: MutableMapping,
        outputs: SequenceClassifierOutputWithPast | None = None,
    ) -> torch.Tensor:
        """Perform a forward pass during evaluation and return the logits.

        This function processes a batch of input data, performs a forward pass using the
        model, and returns the logits for evaluation. It ensures that the "labels"
        component is removed from the batch to avoid computing loss during evaluation.

        Parameters
        ----------
        batch : dict
            A dictionary containing the input data for the model. It must include the
            "labels" key, which will be removed before the forward pass.
        outputs : Any | None, optional
            An optional parameter for additional outputs. Default is None.

        Returns
        -------
        torch.Tensor
            The logits obtained from the forward pass. If the model output is of type
            SequenceClassifierOutputWithPast, the logits are extracted from the output.
            Otherwise, the appropriate tensor is selected based on the output shape.

        Example
        -------
        >>> batch = {
        ...     "input_ids": ...,
        ...     "attention_mask": ...,
        ...     "labels": ...
        ... }
        >>> model = MPTForSequenceClassification(...)
        >>> logits = model.eval_forward(batch)
        >>> print(logits)

        """
        # Pop "labels" component first to avoid computing loss
        self.labels = batch.pop("labels", None)
        assert self.labels is not None
        if outputs is not None:
            # If outputs are provided, use them directly
            output = self._interpret_output(outputs)
        else:
            fwd_outputs = self.forward(batch)
            output = self._interpret_output(fwd_outputs)

        return output
