"""Module implements a DEPToE Causal Language Model for LLM Foundry."""

from __future__ import annotations

import logging
from functools import cached_property
from typing import (
    TYPE_CHECKING,
    Any,
    cast,
)

import torch
from composer.models import HuggingFaceModel
from flash_attn.losses.cross_entropy import (
    CrossEntropyLoss as FusedCrossEntropyLoss,
)
from llmfoundry.metrics import (
    DEFAULT_CAUSAL_LM_EVAL_METRICS,
    DEFAULT_CAUSAL_LM_TRAIN_METRICS,
)
from llmfoundry.registry import models
from llmfoundry.utils.builders import build_metric
from torch import nn

from mt_dao.models.deptoe.configuration_deptoe import DEPToEConfig
from mt_dao.models.deptoe.deptoe_for_causal_lm import DEPToEForCausalLM
from mt_dao.models.ffns.sigma_moe import SigmaMoE
from mt_dao.models.switch_head.switch_head_core import SwitchHead
from mt_dao.models.utils import (
    CROSS_ENTROPY_IGNORE_INDEX,
    compute_loss_from_logits,
    dtensor_safe_check_numel,
)

if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping, MutableMapping

    from transformers.modeling_outputs import (
        CausalLMOutputWithPast,
    )
    from transformers.modeling_utils import (
        PreTrainedModel,
    )
    from transformers.tokenization_utils import (
        PreTrainedTokenizer,
    )
    from transformers.tokenization_utils_fast import (
        PreTrainedTokenizerFast,
    )

    from mt_dao.models.deptoe.deptoe import DEPToEModel

from flwr.common import log


@models.register_class("deptoe_causal_lm")
class DEPToECausalLM(HuggingFaceModel):
    """A 🤗 Transformers compatible DEPToE for causal language modeling."""

    def __init__(
        self,
        tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast | None = None,
        additional_train_metrics: list | None = None,
        loss_fn: str | dict | None = "fused_crossentropy",
        *,
        use_train_metrics: bool | None = True,
        **kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize a DEPToE Causal Language Model.

        Parameters
        ----------
        tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast | None
            The tokenizer to use for the model. If None, the model will not have a
            tokenizer and will not be able to process text inputs.
        additional_train_metrics: list | None
            Additional training metrics to use for the model. If None, no additional
            metrics will be used.
        loss_fn: str | dict | None
            The loss function to use for the model. If None, the default loss function
            will be used. If a string, it must be one of ["fused_crossentropy",
            "torch_crossentropy"]. If a dict, it should contain the configuration for
            the loss function.
        use_train_metrics: bool | None
            Whether to use training metrics. If None, defaults to True.
        **kwargs: Any
            Additional keyword arguments to pass to the model configuration class.

        Raises
        ------
        ValueError
            If the specified `loss_fn` is not recognized.
        ValueError
            If Fused Cross Entropy is not installed and `loss_fn` is set to
            "fused_crossentropy".
        ValueError
            If the specified `loss_fn` is not one of ["fused_crossentropy",
            "torch_crossentropy"].
        ValueError
            If `use_train_metrics` is not a boolean value.

        """
        additional_train_metrics = additional_train_metrics or []

        model = DEPToEForCausalLM(DEPToEConfig(**kwargs))

        train_metric_names = DEFAULT_CAUSAL_LM_TRAIN_METRICS + additional_train_metrics
        train_metrics = (
            [build_metric(metric, {}) for metric in train_metric_names]
            if use_train_metrics
            else []
        )
        eval_metric_names = DEFAULT_CAUSAL_LM_EVAL_METRICS + additional_train_metrics
        eval_metrics = [build_metric(metric, {}) for metric in eval_metric_names]

        super().__init__(
            model=cast("PreTrainedModel", model),
            tokenizer=tokenizer,  # type: ignore[arg-type]
            use_logits=True,
            metrics=train_metrics,
            eval_metrics=eval_metrics,
            shift_labels=model.transformer.shift_labels,
            allow_embedding_resizing=True,
        )

        loss_fn_config = loss_fn
        if loss_fn_config == "fused_crossentropy":
            # TODO(anon): This only works with Flash Attention v2.
            self.loss_fn = FusedCrossEntropyLoss(
                ignore_index=CROSS_ENTROPY_IGNORE_INDEX,
                reduction="none",
            )
        elif loss_fn_config == "torch_crossentropy":
            self.loss_fn = nn.CrossEntropyLoss(
                ignore_index=CROSS_ENTROPY_IGNORE_INDEX,
                reduction="none",
            )
        else:
            msg = f"Specified loss_fn={self.loss_fn} not recognized. `loss_fn` must be"
            msg += " one of [`fused_crossentropy`, `torch_crossentropy`]."
            raise ValueError(msg)
        self.printing_batch_keys = True

    @property
    def model_class(self) -> type[DEPToECausalLM]:
        """Returns the model class for DEPToE underlying model."""
        return DEPToECausalLM

    @property
    def config_class(self) -> type[DEPToEConfig]:
        """Returns the config class for DEPToE underlying model."""
        return DEPToEConfig

    def forward(self, batch: MutableMapping) -> CausalLMOutputWithPast:
        """Forward pass of the DEPToE Causal Language Model.

        Parameters
        ----------
        batch : MutableMapping
            A batch of data containing the input IDs, attention mask, sequence IDs,
            inputs embeddings, and position IDs.

        Returns
        -------
        CausalLMOutputWithPast
            The output of the model, which includes the logits and other information.

        """
        if self.printing_batch_keys:
            log(
                logging.INFO,
                "Batch keys: %s",
                ", ".join(batch.keys()),
            )
            #  Batch keys: input_ids, labels, total_tokens, loss_generating_tokens
            self.printing_batch_keys = False

        return cast("DEPToEForCausalLM", self.model)(
            input_ids=batch.get("input_ids", None),
            attention_mask=batch.get("attention_mask", None),
            labels=batch.get("labels", None),
            return_dict=batch.get("return_dict", None),
            output_attentions=batch.get("output_attentions", None),
            use_cache=batch.get("use_cache", None),
        )

    def loss(
        self,
        outputs: CausalLMOutputWithPast,
        batch: Mapping,
    ) -> dict | torch.Tensor:
        """Compute the loss from the model outputs and the batch.

        Parameters
        ----------
        outputs : CausalLMOutputWithPast
            The outputs of the model, which includes the logits and other information.
        batch : Mapping
            The batch of data, which should contain the labels for computing the loss.

        Returns
        -------
        dict | torch.Tensor
            The computed loss. If the model uses MegaBlocks MoE, it returns a dictionary
            containing the total loss, the loss, and the load balancing loss. Otherwise,
            it returns the loss as a tensor.

        """
        loss = compute_loss_from_logits(
            outputs,
            batch["labels"],
            self.loss_fn,
            shift_labels=self.shift_labels,
        )

        # Collect regularizaiton losses across the layers
        if cast("DEPToEForCausalLM", self.model).transformer.n_repeats > 1:
            for layer in self.children():
                if isinstance(layer, SigmaMoE):
                    loss += (
                        cast("DEPToEForCausalLM", self.model).entropy_reg
                        * layer.get_reg_loss()
                    )
                elif isinstance(layer, SwitchHead):
                    loss += (
                        cast("DEPToEForCausalLM", self.model).att_entropy_reg
                        * layer.get_reg_loss()
                    )

        return loss

    # TODO(anon): Check if the following methods are doing what we expect them to do
    @cached_property
    def n_total_params(self) -> int:
        """Gets the total number of parameters in the model.

        Returns
        -------
        int
            The total number of parameters in the model, including both trainable and
            non-trainable parameters.

        """
        return sum(
            p.numel() for p in cast("DEPToEForCausalLM", self.model).parameters()
        )

    @cached_property
    def n_trainable_params(self) -> int:
        """Gets the total trainable number of parameters in the model.

        Returns
        -------
        int
            The total trainable number of parameters in the model.

        """
        return sum(
            p.numel()
            for p in cast("DEPToEForCausalLM", self.model).parameters()
            if p.requires_grad
        )

    @cached_property
    def n_active_params(self) -> int:
        """Gets the total number of active parameters in the model.

        Returns
        -------
        int
            The total number of active parameters in the model, which are the parameters
            that are trainable and not frozen.

        """
        params = sum(p.numel() for p in self.parameters())
        if not cast(
            "DEPToEForCausalLM",
            self.model,
        ).transformer.config.tie_word_embeddings:
            # Embedding layers are lookup tables, therefore are not counted in the FLOP
            # computation
            params -= dtensor_safe_check_numel(
                cast("DEPToEForCausalLM", self.model).transformer.wte.weight,
            )
        return params

    @cached_property
    def n_non_embedding_params(self) -> int:
        """Count the number of non-embedding parameters in the model.

        If the model uses tied word embeddings, this will exclude the embedding layer
        from the count. Otherwise, it will include all parameters except the embedding
        layer and the unembedding language modeling head.

        Returns
        -------
        int
            The number of non-embedding parameters in the model. This is useful for
            calculating the number of parameters that are not related to the embedding
            layer, which is often needed for FLOP calculations or model size
            estimations.

        """
        # Total number of parameters minus the embedding layer and the unembedding
        # language modeling head, if they exist.
        params = sum(p.numel() for p in self.parameters())
        params -= dtensor_safe_check_numel(
            cast("DEPToEForCausalLM", self.model).transformer.wte.weight,
        )
        if hasattr(self.model, "lm_head") and self.model.lm_head is not None:
            params -= dtensor_safe_check_numel(
                cast("torch.nn.Linear", self.model.lm_head).weight,
            )
        return params

    def flops_per_batch(self, batch: Mapping) -> int:
        """Compute the number of FLOPs per batch.

        Parameters
        ----------
        batch : Mapping
            The batch of data for which to compute the FLOPs. It should contain the
            `input_ids` key, which is a tensor of shape (batch_size, sequence_length).

        Returns
        -------
        int
            The number of FLOPs per batch. This is computed as the sum of the FLOPs
            for the parameters and the attention mechanism, multiplied by the batch
            size.

        """
        # NOTE: this computation does not take into account padding, and assumes
        # that the dataset has been constructed without padding. Additionally, we
        # assume the backward pass is approximately 2x the forward pass

        bs, msl = cast("torch.Tensor", batch["input_ids"]).shape[0:2]
        params = self.n_active_params
        params_flops_per_token = 2 * params
        params_flops_per_seq = params_flops_per_token * msl
        attn_flops_per_seq = self.get_attention_flops(msl)
        return (params_flops_per_seq + attn_flops_per_seq) * 3 * bs

    def get_attention_flops(self, msl: int) -> int:
        """Compute the number of FLOPs for the attention mechanism.

        Parameters
        ----------
        msl : int
            The maximum sequence length for the attention mechanism.

        Returns
        -------
        int
            The number of FLOPs for the attention mechanism, which is computed as the
            product of the number of layers, the number of heads, the model dimension,
            and the square of the maximum sequence length.

        """
        return (
            cast("DEPToEModel", self.model).config.n_layers
            * 2
            * 2
            * (cast("DEPToEModel", self.model).config.d_model * (msl**2))
        )

    def get_optimizer_param_groups(
        self,
        optimizer_config: dict[str, Any],
    ) -> tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]:
        """Extract the parameter groups for the optimizer when using MUP.

        Delegated to the underlying model's method.

        Parameters
        ----------
        optimizer_config : dict[str, Any]
            The configuration for the optimizer, which may include parameters like
            learning rate, weight decay, etc.

        Returns
        -------
        tuple[Iterable[torch.Tensor] | Iterable[dict[str, Any]], dict[str, Any]]
            A tuple containing the parameter groups and the optimizer configuration.

        """
        return cast(
            "DEPToEForCausalLM",
            self.model,
        ).get_optimizer_param_groups(optimizer_config)
