"""Module implements MPT Model for Token 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 CrossEntropyLoss
from torchmetrics import Metric
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,
    TokenClassifierOutput,
)

if TYPE_CHECKING:
    from llmfoundry.models import MPTForCausalLM

logger = logging.getLogger(__name__)


class MPTForTokenClassification(HuggingFaceModel):
    """MPT model for token classification tasks."""

    def __init__(  # noqa: PLR0913, PLR0917
        self,
        model: ComposerMPTCausalLM,
        num_labels: int,
        hidden_size: int,
        classifier_dropout: float | None = None,
        hidden_dropout: float | None = None,
        train_metrics: list[Metric] | None = None,
        eval_metrics: list[Metric] | None = None,
    ) -> None:
        """Initialize the MPTForTokenClassification model.

        Parameters
        ----------
        model : ComposerMPTCausalLM
            The transformer model to be used for token classification.
        num_labels : int
            The number of labels for token classification.
        hidden_size : int
            The size of the hidden layer in the transformer model.
        classifier_dropout : float | None, optional
            The dropout rate for the classifier layer (default is None).
        hidden_dropout : float | None, optional
            The dropout rate for the hidden layer (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

        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
        if classifier_dropout is not None:
            dropout = classifier_dropout
        elif hidden_dropout is not None:
            dropout = hidden_dropout
        else:
            dropout = 0.1
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(
        self,
        batch: MutableMapping,
    ) -> (
        tuple[
            torch.Tensor | None,
            torch.Tensor,
            tuple[torch.FloatTensor] | None,
            tuple[torch.FloatTensor] | None,
        ]
        | TokenClassifierOutput
    ):
        """Perform a forward pass for token classification.

        Parameters
        ----------
        batch : MutableMapping
            A dictionary containing the input data for the model, including:
            - "input_ids": Tensor of input token IDs.
            - "attention_mask": Tensor indicating which tokens are padding.
            - "labels": Optional tensor of labels for token classification.
            - "output_hidden_states": Optional flag to return hidden states.
            - "return_dict": Optional flag to return a TokenClassifierOutput object.

        Returns
        -------
        TokenClassifierOutput | tuple
            An object containing the logits, loss (if labels are provided),
            hidden states, and attentions. If `return_dict` is False, returns a tuple
            containing the logits, hidden states, and attentions.

        """
        labels = batch.pop("labels", None)
        output_hidden_states = batch.get("output_hidden_states", None)
        return_dict = batch.get("return_dict", None)

        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,
        )

        hidden_states = transformer_outputs[0]
        hidden_states = self.dropout(hidden_states)
        logits = cast("torch.Tensor", self.classifier(hidden_states))

        loss = None
        if labels is not None:
            # Move labels to correct device to enable model parallelism
            labels = labels.to(logits.device)
            batch_size, seq_length = labels.shape
            loss_fct = CrossEntropyLoss()
            loss = cast(
                "torch.Tensor",
                loss_fct(
                    logits.view(batch_size * seq_length, self.num_labels),
                    labels.view(batch_size * seq_length),
                ),
            )

        if not return_dict:
            output = cast(
                """tuple[
                    torch.Tensor,
                    tuple[torch.FloatTensor] | None,
                    tuple[torch.FloatTensor] | None
                ]""",
                (
                    logits,
                    transformer_outputs.hidden_states,
                    transformer_outputs.attentions,
                ),
            )
            return (loss, *output)

        return TokenClassifierOutput(
            loss=cast("torch.FloatTensor", loss) if loss is not None else None,
            logits=cast("torch.FloatTensor", logits),
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )
