"""Module implements MPT Model for Question Answering."""
import logging
from collections.abc import MutableMapping
from typing import TYPE_CHECKING, Any, cast

import torch
from composer.models import HuggingFaceModel
from llmfoundry import ComposerMPTCausalLM
from torch import FloatTensor, nn
from torch.nn import CrossEntropyLoss
from torchmetrics import Metric
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,
    QuestionAnsweringModelOutput,
)

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

logger = logging.getLogger(__name__)


class MPTForQuestionAnswering(HuggingFaceModel):
    """MPT Model for Question Answering.

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

    Attributes
    ----------
    transformer : ComposerMPTCausalLM
        The transformer model used for question answering.
    qa_outputs : nn.Linear
        A linear layer to compute the start and end logits for the answer span.

    Methods
    -------
    __init__(model, hidden_size, train_metrics=None, eval_metrics=None)
        Initialize the MPTForQuestionAnswering model.
    forward(batch) -> tuple[torch.Tensor] | QuestionAnsweringModelOutput
        Perform a forward pass on the input batch and return the question answering
        output.

    Parameters
    ----------
    model : ComposerMPTCausalLM
        The transformer model to be used for question answering.
    hidden_size : int
        The size of the hidden layer in the transformer model.
    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(...)
    >>> qa_model = MPTForQuestionAnswering(
    ...     model=model,
    ...     hidden_size=768,
    ...     train_metrics=[...],
    ...     eval_metrics=[...]
    ... )
    >>> batch = {
    ...     "input_ids": ...,
    ...     "attention_mask": ...,
    ...     "start_positions": ...,
    ...     "end_positions": ...
    ... }
    >>> output = qa_model.forward(batch)
    >>> print(output)

    """

    def __init__(
        self,
        model: ComposerMPTCausalLM,
        hidden_size: int,
        dropout_rate: float = 0.1,
        train_metrics: list[Metric] | None = None,
        eval_metrics: list[Metric] | None = None,
    ) -> None:
        """Initialize the MPTForQuestionAnswering model.

        Parameters
        ----------
        model : ComposerMPTCausalLM
            The transformer model to be used for question answering.
        hidden_size : int
            The size of the hidden layer in the transformer model.
        dropout_rate : float, optional
            The dropout rate for the model (default is 0.1).
        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).

        """
        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.dropout = nn.Dropout(dropout_rate)
        self.qa_outputs = nn.Linear(hidden_size, 2, bias=True)

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

        This function processes a batch of input data, performs a forward pass using the
        transformer model, and returns the start and end logits for the answer span. It
        supports different return types based on the `return_dict` parameter and
        computes the loss if start and end positions 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.
            - "attention_mask" (torch.Tensor): The attention mask for the input data.
            - "start_positions" (torch.Tensor): The start positions of the answer spans.
            - "end_positions" (torch.Tensor): The end positions of the answer spans.
            - "return_dict" (bool): Whether to return a dictionary of outputs.

        Returns
        -------
        tuple[torch.Tensor, ...] | QuestionAnsweringModelOutput
            The question answering output. If `return_dict` is True, a
            QuestionAnsweringModelOutput object is returned. Otherwise, a tuple
            containing the start logits, end logits, and additional outputs is returned.
            If start and end positions are provided, the loss is included in the output.

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

        """
        example_ids = batch.pop("example_ids", None)
        answers = batch.pop("answers", None)
        start_positions = batch.pop("start_positions", None)
        end_positions = batch.pop("end_positions", 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 = self.dropout(
            transformer_outputs.hidden_states[-1],
        )  # (bs, max_query_len, dim)
        logits: torch.Tensor = self.qa_outputs(hidden_states)  # (bs, max_query_len, 2)
        start_logits_, end_logits_ = cast(
            "tuple[torch.Tensor, torch.Tensor]",
            logits.split(1, dim=-1),
        )
        # The shape is (bs, max_query_len)
        start_logits = start_logits_.squeeze(-1).contiguous()
        # The shape is (bs, max_query_len)
        end_logits = end_logits_.squeeze(-1).contiguous()

        total_loss: torch.Tensor | None = None
        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, split add a dimension
            if len(start_positions.size()) > 1:
                start_positions = start_positions.squeeze(-1)
            if len(end_positions.size()) > 1:
                end_positions = end_positions.squeeze(-1)
            # Sometimes the start/end positions are outside our model inputs, we ignore
            # these terms
            ignored_index = start_logits.size(1)
            start_positions = start_positions.clamp(0, ignored_index)
            end_positions = end_positions.clamp(0, ignored_index)

            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
            start_loss: torch.Tensor = loss_fct(start_logits, start_positions)
            end_loss: torch.Tensor = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2

        if not return_dict:
            output = (
                start_logits,
                end_logits,
                *cast("tuple[torch.Tensor, ...]", transformer_outputs[2:]),
            )
            return ((total_loss, *output)) if total_loss is not None else output

        # NOTE: We need to put the "answers" back in to propagate it to the output
        batch["answers"] = answers

        # NOTE: We need to put the "example_ids" back in to propagate it to the output
        batch["example_ids"] = example_ids

        return QuestionAnsweringModelOutput(
            loss=(
                cast("torch.FloatTensor", total_loss)
                if total_loss is not None
                else None
            ),
            start_logits=cast("torch.FloatTensor", start_logits),
            end_logits=cast("torch.FloatTensor", end_logits),
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

    def eval_forward(
        self,
        batch: MutableMapping,
        outputs: QuestionAnsweringModelOutput | None = None,
    ) -> tuple[FloatTensor, FloatTensor, MutableMapping[Any, Any]]:
        """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
            QuestionAnsweringModelOutput, 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 = MPTForQuestionAnswering(...)
        >>> logits = model.eval_forward(batch)
        >>> print(logits)

        """
        # Pop "labels" component first to avoid computing loss
        self.labels = batch.pop("labels", None)
        if outputs is not None and type(outputs) is QuestionAnsweringModelOutput:
            assert (
                outputs.start_logits is not None
            ), "outputs.start_logits must not be None"
            assert (
                outputs.end_logits is not None
            ), "outputs.end_logits must not be None"
            return outputs.start_logits, outputs.end_logits, batch
        output = self.forward(batch)

        assert type(output) is QuestionAnsweringModelOutput
        assert (
            output.start_logits is not None
        ), "output.start_logits must not be None"
        assert (
            output.end_logits is not None
        ), "output.end_logits must not be None"

        return output.start_logits, output.end_logits, batch
