"""Module implements the backbone DEPToE model for LLM Foundry.

For MuP/CompleteP configuration see:

References
----------
- CompleteP: https://arxiv.org/abs/2505.01618
- MuP: https://arxiv.org/abs/2203.03466
- Practitioner's Guide to MuP: https://blog.eleuther.ai/mutransfer/
- Original Code: https://github.com/EleutherAI/nanoGPT-mup
- Our Code: https://github.com/anon-anon/nanoGPT-mt_dao

"""

from __future__ import annotations

import copy
import types
from logging import DEBUG, INFO
from typing import (
    Any,
    ClassVar,
    cast,
)

import torch
import torch.nn.functional as F  # noqa: N812
from composer.utils import dist
from flwr.common import log
from flwr.common import log as flwr_log
from llmfoundry.models.layers.custom_embedding import SharedEmbedding
from llmfoundry.models.layers.layer_builders import build_norm
from llmfoundry.models.utils.act_ckpt import (
    pass_on_block_idx,
)
from torch import nn
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
)
from transformers.modeling_utils import (
    PreTrainedModel,
)

from mt_dao.models.deptoe.configuration_deptoe import DEPToEConfig
from mt_dao.models.deptoe.deptoe_block import DEPToEBlock
from mt_dao.models.deptoe.mup_config import MuPConfig
from mt_dao.models.ffns.sigma_moe import SigmaMoE
from mt_dao.models.rope.rope_encoding import RotaryPositionalEncoding
from mt_dao.models.switch_head.switch_head_core import SwitchHead

CROSS_ENTROPY_IGNORE_INDEX = -100


class DEPToEPreTrainedModel(PreTrainedModel):
    """Base class for DEPToE PreTrainedModel.

    This class is used to define the configuration and base model prefix for DEPToE
    models. It inherits from PreTrainedModel and sets the configuration class to
    DEPToEConfig. It also defines the base model prefix as "model" and specifies that
    the no-split modules are "DEPToEBlock".

    """

    config_class = DEPToEConfig  # type: ignore[assignment]
    base_model_prefix = "model"  # type: ignore[misc]
    _no_split_modules: ClassVar[list[str]] = ["DEPToEBlock"]  # type: ignore[assignment,misc]


class DEPToEModel(DEPToEPreTrainedModel):
    """DEPToE Model class."""

    def __init__(self, config: DEPToEConfig) -> None:
        """Initialize the DEPToE model.

        Parameters
        ----------
        config : DEPToEConfig
            The configuration for the DEPToE model.

        Raises
        ------
            NotImplementedError
                If the requested norm type is not implemented.

        """
        config.validate_config()
        super().__init__(config)
        self.config: DEPToEConfig = cast("DEPToEConfig", self.config)  # type: ignore[reportUnnecessaryCast]
        self.mup_config = MuPConfig(**self.config.mup_config)  # type: ignore[reportArgumentType]

        # TODO(anon): Check this out
        if config.init_device == "mixed":
            if dist.get_local_rank() == 0:
                config.init_device = "cpu"
            else:
                config.init_device = "meta"

        # TODO(anon): Check this out
        # CogView (https://arxiv.org/abs/2105.13290) and GLM-130B (https://arxiv.org/abs/2210.02414)
        # both report this helping with stabilizing training
        self.embedding_fraction = config.embedding_fraction

        self.wte = SharedEmbedding(
            config.vocab_size,
            config.d_model,
            padding_idx=config.pad_token_id,
            device=config.init_device,
        )
        self.embedding_dropout = nn.Dropout(config.embedding_dropout)
        # TODO(anon): Check this out
        self.shift_labels = True

        # Initialize the Rotary Positional Encoding (RoPE)
        self.pe = RotaryPositionalEncoding(
            attn_heads_projection_size=config.attn_heads_projection_size,
            base=config.rope_base,
            scale_base=config.attn_xpos_scale_base,
            interleaved=config.rope_interleaved,
            use_torch_apply_rotary_emb=config.use_torch_apply_rotary_emb,
        )

        self.blocks = self.construct_blocks(config=config)
        self.config.deptoe_block_config = None  # type: ignore[reportAttributeAccessIssue,assignment]

        # Set the number of UT repeats
        self.n_repeats = config.n_repeats

        # Set the PeriLN flags
        self.use_peri_ln = config.use_peri_ln
        self.use_embedding_peri_ln = config.use_embedding_peri_ln

        # Add the embedding PeriLN layer
        # NOTE: With or without PeriLN we do have an input LayerNorm
        if self.use_peri_ln and self.use_embedding_peri_ln:
            self.embedding_peri_ln = torch.nn.LayerNorm(
                config.d_model,
            )

        # TODO(anon): Check this out
        # Tag all modules in the transformer blocks with the corresponding block_idx and
        # max_block_idx
        for i, block in enumerate(self.blocks):
            assert isinstance(block, self.block_class), (
                f"Expected block to be of type {self.block_class},"
                f" but got {type(block)}."
            )
            block.block_idx = torch.Tensor(i)
            block.max_block_idx = torch.Tensor(config.n_layers - 1)
            pass_on_block_idx(block)

        # NOTE: this may or may not have a bias
        self.norm_f = build_norm(
            name=config.out_norm_type.lower(),
            normalized_shape=config.d_model,
            eps=config.out_norm_eps,
            device=config.init_device,
        )

        # NOTE: code to shut down all biases if present
        # stolen from MPT
        if config.no_bias:
            for module in self.modules():
                if hasattr(module, "bias") and isinstance(module.bias, nn.Parameter):
                    module.register_parameter("bias", None)

                # For transformer engine
                if hasattr(module, "use_bias") and module.use_bias is True:
                    module.use_bias = False  # type: ignore[reportArgumentType]

        if config.init_device != "meta":
            log(
                INFO,
                'We recommend using config.init_device="meta" with Composer + FSDP for'
                " faster initialization.",
            )
            # TODO(anon): Make sure to format the strings in this log properly.
            log(
                INFO,
                "Applying parameter initialization function %s to children of %s: %s.",
                self.reset_parameters.__name__,
                self.__class__.__name__,
                [str(child) for child in self.children()],
            )
            # NOTE: If the device is "meta" (only for FSDP), all module's parameters
            # are initialized through their respective module's `reset_parameters`
            # method.
            self.reset_parameters()

    def reset_parameters(self) -> None:
        """Reset the parameters of the model.

        This method initializes the parameters of the model using a scale factor
        based on the number of repeats and the number of blocks. It also initializes
        the weights and biases of the LayerNorm layers to ones and zeros, respectively.
        The scale factor is calculated as the square root of 2 divided by the product
        of the number of repeats and the number of blocks.
        This method is called without gradients to ensure that the parameters are
        initialized correctly without affecting the gradient computation.

        """
        with torch.no_grad():
            self.init_trainable_embeddings(
                self.wte,
                self.config.init_config,
                mup_enabled=self.mup_config.mup_enabled,
            )
            for layer in self.modules():
                if isinstance(layer, (SwitchHead | SigmaMoE)):
                    layer.reset_parameters()
                elif isinstance(layer, torch.nn.LayerNorm):
                    torch.nn.init.ones_(layer.weight)
                    if layer.bias is not None:  # type: ignore[reportUnnecessaryComparison]
                        torch.nn.init.zeros_(layer.bias)

    @property
    def block_class(self) -> type[DEPToEBlock]:
        """Return the class of the transformer block."""
        return DEPToEBlock

    def construct_blocks(self, config: DEPToEConfig) -> nn.ModuleList:
        """Construct the transformer blocks for the model.

        Parameters
        ----------
        config : DEPToEConfig
            The configuration for the DEPToE model.

        Returns
        -------
        nn.ModuleList
            A list of transformer blocks initialized with the given configuration.

        """
        deptoe_block_config = config.deptoe_block_config
        deptoe_block_config.device = config.init_device

        return nn.ModuleList(
            [
                self.block_class(
                    config=copy.deepcopy(deptoe_block_config),
                    rope_module=self.pe,
                )
                for _ in range(config.n_layers)
            ],
        )

    # TODO(anon): Check this out
    @staticmethod
    def _resolve_reuse_kv_layer_idx(
        overrides_definition: dict[str, Any],
        model_modules_order_expanded: list[str],
        b_idx: int,
        override_config: dict[str, Any],
        reuse_kv_layer_idx_dict: dict[int, int],
    ) -> int:
        override_attn_config = override_config["attn_config"]
        if override_attn_config["reuse_kv_layer_idx"] >= 0:
            reuse_kv_layer_idx = override_attn_config["reuse_kv_layer_idx"]
            msg = "The relative index of kv layer to reuse,"
            msg += f" override_attn_config['reuse_kv_layer_idx']={reuse_kv_layer_idx},"
            msg += " should be negative."
            raise ValueError(msg)
        reuse_kv_layer_idx = b_idx + override_attn_config["reuse_kv_layer_idx"]
        if reuse_kv_layer_idx < 0:
            msg = f"The absolute index of kv layer to reuse, {reuse_kv_layer_idx}"
            msg += " should be non-negative."
            raise ValueError(msg)
        if reuse_kv_layer_idx in reuse_kv_layer_idx_dict:
            reuse_kv_layer_idx = reuse_kv_layer_idx_dict[reuse_kv_layer_idx]
        reuse_kv_layer_idx_dict[b_idx] = reuse_kv_layer_idx

        parent_layer_name = model_modules_order_expanded[reuse_kv_layer_idx]
        parent_config = (
            {}
            if parent_layer_name == "default"
            else copy.deepcopy(
                overrides_definition[parent_layer_name],
            )
        )
        if "attn_config" not in parent_config:
            parent_config["attn_config"] = {}
        parent_config["attn_config"]["reuse_kv_layer_idx"] = override_config[
            "attn_config"
        ]["reuse_kv_layer_idx"]

        if override_config != parent_config and not override_config.get(
            "allow_mismatch",
        ):
            msg = "For reusing the kv cache of a previous layer, the previous layer"
            msg += " should match the block config as the current layer."
            raise ValueError(msg)

        return reuse_kv_layer_idx

    @staticmethod
    def _get_modules_order_expanded(order: list[dict[str, Any]]) -> list[str]:
        model_modules_order_expanded = []
        for item in order:
            repeat = item.get("repeat", 1)
            if ("name" in item) == ("order" in item):
                msg = (
                    "Exactly one of `order` or `name` must be specified for each"
                    " block override."
                )
                raise ValueError(msg)

            if "name" in item:
                model_modules_order_expanded.extend([item["name"]] * repeat)
            else:
                model_modules_order_expanded.extend(
                    DEPToEModel._get_modules_order_expanded(item["order"]) * repeat,
                )

        return model_modules_order_expanded

    # TODO(anon): Add DEPT embeddings to the following methods
    def get_input_embeddings(self) -> SharedEmbedding | nn.Embedding:
        """Get the input embeddings of the model.

        Returns
        -------
        SharedEmbedding | nn.Embedding: The input embeddings of the model.

        """
        return self.wte

    def set_input_embeddings(  # type: ignore[reportIncompatibleMethodOverride]
        self,
        value: SharedEmbedding | nn.Embedding,
    ) -> None:
        """Set the input embeddings of the model.

        Parameters
        ----------
        value : SharedEmbedding | nn.Embedding
            The input embeddings to set.

        """
        # NOTE: If we use a different tokenizer from gpt2
        # we need to make sure we are forward and backward
        # compatible with the SharedEmbedding of MosaicML
        # Which requires the `unembed` parameter to work
        if not isinstance(value, SharedEmbedding):
            original_forward = value.forward

            def new_forward(
                self: SharedEmbedding | nn.Embedding,
                input_tensor: torch.Tensor,
                *,
                unembed: bool = False,
            ) -> torch.Tensor:
                if unembed:
                    return F.linear(input_tensor, self.weight)
                return original_forward(input_tensor)

            value.forward = types.MethodType(new_forward, value)
        self.wte = value

    def _forward_checks(
        self,
        attention_mask: torch.ByteTensor | None = None,
        input_ids: torch.LongTensor | None = None,
        *,
        return_dict: bool | None = None,
        output_attentions: bool | None = None,
        use_cache: bool | None = None,
    ) -> None:
        return_dict = (
            return_dict if return_dict is not None else self.config.return_dict
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if use_cache:
            msg = "use_cache=True is not implemented yet for DEPToEModel."
            raise NotImplementedError(msg)

        # TODO(anon): These args are passed in by keyword in huggingface's generate
        # function: https://github.com/huggingface/transformers/blob/68287689f2f0d8b7063c400230b3766987abf18d/src/transformers/generation/utils.py#L2201-L2206
        # but have not yet been fully implemented in DEPToEModel
        if not return_dict:
            msg = "return_dict=False is not implemented yet for DEPToEModel."
            raise NotImplementedError(msg)
        if output_attentions and self.attn_impl != "torch":
            msg = (
                "output_attentions is not implemented for DEPToEModel when using"
                " attn_impl `flash`."
            )
            raise NotImplementedError(msg)

        if (
            self.training
            and attention_mask is not None
            and attention_mask[:, 0].sum() != attention_mask.shape[0]
        ):
            msg = "DEPToEModel does not support training with left padding."
            raise NotImplementedError(msg)

        if input_ids is None:
            msg = "You must specify input_ids."
            raise ValueError(msg)

    def apply_embedding_layer(self, input_ids: torch.LongTensor | None) -> torch.Tensor:
        """Apply the embedding layer to the input IDs.

        Parameters
        ----------
        input_ids : torch.LongTensor | None
            The input IDs to be embedded.

        Returns
        -------
        torch.Tensor
            The embedded input tensor.

        """
        x = cast("torch.Tensor", self.wte(input_ids))

        sequence_length = self.get_sequence_length(x)

        assert sequence_length <= self.config.max_seq_len, (
            f"Cannot forward input with seq_len={sequence_length}, this model only"
            f" supports seq_len<={self.config.max_seq_len}"
        )

        if self.embedding_fraction == 1:
            x = self.embedding_dropout(x)
        else:
            # NOTE: this implementation is proposed on page 7 of the GLM-130B paper https://arxiv.org/abs/2210.02414
            x_shrunk = (x * self.embedding_fraction) + (
                x.detach() * (1 - self.embedding_fraction)
            )
            assert isinstance(self.embedding_dropout, nn.Module)
            x = self.embedding_dropout(x_shrunk)

        # Layer norm here as it's done in the Peri-LN paper: https://arxiv.org/html/2502.02732v1
        if self.use_peri_ln:
            x = self.embedding_peri_ln(x)

        return x

    def forward(  # noqa: PLR0913
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.ByteTensor | None = None,
        *,
        return_dict: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        use_cache: bool | None = None,
    ) -> BaseModelOutputWithPast:
        """Forward pass of the DEPToEModel.

        Parameters
        ----------
        input_ids : torch.LongTensor | None, optional
            The input IDs of shape (batch_size, seq_len), by default None.
        attention_mask : torch.ByteTensor | None, optional
            The attention mask of shape (batch_size, seq_len), by default None.
        return_dict : bool | None, optional
            Whether to return a dictionary of outputs, by default None.
        output_attentions : bool | None, optional
            Whether to return attention weights, by default None.
        output_hidden_states : bool | None, optional
            Whether to return hidden states, by default None.
        use_cache : bool | None, optional
            Whether to use caching for key-value pairs, by default None.

        Returns
        -------
        BaseModelOutputWithPast
            The output of the model containing the last hidden state, past key values,
            hidden states, and attentions if requested.


        """
        self._forward_checks(
            attention_mask=attention_mask,
            input_ids=input_ids,
            return_dict=return_dict,
            output_attentions=output_attentions,
            use_cache=use_cache,
        )

        x = self.apply_embedding_layer(input_ids)

        # NOTE: Input scaling enables more control
        # enabling it to match specific SP parametrizations
        # usually set to 1.0
        if self.mup_config.mup_enabled:
            x = x * self.mup_config.mup_input_alpha  # noqa: PLR6104

        all_hidden_states: tuple[torch.FloatTensor, ...] | None = (
            cast("tuple[torch.FloatTensor, ...]", (x,))
            if output_hidden_states
            else None
        )
        all_self_attns: tuple[torch.FloatTensor, ...] | None = (
            () if output_attentions else None
        )

        for _ in range(self.n_repeats):
            for _b_idx, (block) in enumerate(cast("list[DEPToEBlock]", self.blocks)):
                x, _present, attn_weights = cast(
                    """tuple[
                        torch.Tensor,
                        torch.Tensor | None,
                        tuple[torch.Tensor, torch.Tensor] | None
                    ]""",
                    block(
                        x=x,
                        past_key_value=None,
                        attention_mask=attention_mask,
                        output_attentions=output_attentions,
                    ),
                )
                # NOTE: At the first iteration we add hidden states from the embedding
                if output_hidden_states:
                    assert all_hidden_states is not None
                    all_hidden_states += (cast("torch.FloatTensor", x),)

                if output_attentions:
                    assert all_self_attns is not None
                    all_self_attns += (cast("torch.FloatTensor", attn_weights),)

        x = cast("torch.Tensor", self.norm_f(x))

        # NOTE: the output logits must be divided by the width multiplier
        # to maintain the scale of the activations constant across widths
        # the output_alpha works similarly to the input_alpha
        # allowing MuP to match specific SP parametrizations
        # Nanogpt has the scaling after the norm prior to unembedding
        if self.mup_config.mup_enabled:
            x = (
                x
                * self.mup_config.mup_output_alpha
                / self.mup_config.mup_width_multiplier
            )

        # NOTE: Add the states after the final layer norm
        if output_hidden_states:
            assert all_hidden_states is not None
            all_hidden_states += (cast("torch.FloatTensor", x),)

        return BaseModelOutputWithPast(
            last_hidden_state=cast("torch.FloatTensor", x),
            past_key_values=None,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

    @staticmethod
    def get_sequence_length(x: torch.Tensor) -> int:
        """Return the sequence length.

        Parameters
        ----------
        x : torch.Tensor
            The input tensor.

        Returns
        -------
        int: The sequence length of the input tensor.

        """
        return x.size(1)

    # FSDP Wrap function
    @staticmethod
    def fsdp_wrap_fn(module: nn.Module) -> bool:
        """FSDP wrap function for DEPToEModel.

        Parameters
        ----------
        module : nn.Module
            The module to check for FSDP wrapping.

        Returns
        -------
        bool: True if the module is a DEPToEBlock, False otherwise.

        """
        # FSDP Wrap function for MPT Models
        if hasattr(module, "_fsdp_kwargs_dict"):
            log(
                DEBUG,
                "DEPToEModel fsdp_wrap_fn called to %s with _fsdp_kwargs_dict=%s",
                module,
                module._fsdp_kwargs_dict,  # noqa: SLF001
            )
            return bool(module._fsdp_kwargs_dict)  # noqa: SLF001
        log(DEBUG, "DEPToEModel fsdp_wrap_fn called to %s", module)
        return isinstance(module, DEPToEBlock)

    @staticmethod
    def activation_checkpointing_fn(module: nn.Module) -> bool:
        """Check if the module is a DEPToEBlock for activation checkpointing.

        Parameters
        ----------
        module : nn.Module
            The module to check for activation checkpointing.

        Returns
        -------
        bool: True if the module is a DEPToEBlock, False otherwise.

        """
        return isinstance(module, DEPToEBlock)

    @staticmethod
    def init_trainable_embeddings(
        module: nn.Embedding | SharedEmbedding | nn.Linear,
        init_config: dict[str, Any],
        *,
        mup_enabled: float = False,
    ) -> None:
        """Initialize the weights of the embedding layers.

        This method initializes the weights of the embedding layers, including
        `SharedEmbedding`, `torch.nn.Embedding`, and `nn.Linear` layers with no bias.
        It uses Kaiming normal initialization for the weights of these layers if
        `mup_enabled` is False. If `mup_enabled` is True, it uses a Standard Gaussian
        with a standard deviation specified in the configuration.

        Parameters
        ----------
        module : nn.Module
            The module to initialize. It can be an instance of `SharedEmbedding`,
            `torch.nn.Embedding`, or `nn.Linear` with no bias.
        init_config : dict[str, Any]
            The initialization configuration dictionary containing the initialization
            parameters.
        mup_enabled : bool, optional
            Whether MuP (Multi-Path) is enabled. If True, the initialization will use
            a Standard Gaussian distribution with a specified standard deviation.


        """
        if not mup_enabled:
            torch.nn.init.kaiming_normal_(
                module.weight,
                mode="fan_in",
                nonlinearity="linear",
            )
        else:
            # NOTE: Tensor program 5 specifies Gaussian initialization
            # arguing that all other initializations converge to a Gaussian one
            # in the infinite-width limit.
            std = init_config.get("emb_init_std", 1.0)
            std = std if std is not None else init_config.get("init_std", 1.0)

            # NOTE: 1.0 is the default of the normal_ distribution
            std = std if std is not None else 1.0
            torch.nn.init.normal_(
                module.weight,
                mean=0.0,
                std=std,
            )
            flwr_log(
                DEBUG,
                "Initialized %s with std=%s",
                module.__class__.__name__,
                std,
            )
