"""MoEUT-style Blocks used for the DEPToE Model."""

from typing import TYPE_CHECKING, cast

import torch
from torch import nn

from mt_dao.models.deptoe.mup_config import MuPConfig
from mt_dao.models.ffns.sigma_moe import SigmaMoE, SigmaMoEConfig
from mt_dao.models.rope.rope_encoding import RotaryPositionalEncoding
from mt_dao.models.switch_head.switch_head_core import SwitchHead, SwitchHeadConfig

if TYPE_CHECKING:
    from mt_dao.models.switch_head.utils import KVCache

__all__ = [
    "DEPToEBlock",
]


class DEPToEBlockConfig:
    def __init__(  # noqa: PLR0913, PLR0917
        self,
        d_model: int,
        n_heads: int,
        n_repeats: int = 1,
        n_blocks: int = 1,
        attn_heads_experts: int = 1,
        attn_heads_dropout: float = 0.0,
        attn_heads_projection_size: int = 64,
        attn_heads_experts_dropout: float = 0.0,
        attn_heads_n_active_experts: int = 1,
        rope_base: int = 10000,
        ff_n_experts: int = 1,
        ff_expert_size: int = 64,
        ff_n_active_experts: int = 1,
        ff_expert_dropout: float = 0.0,
        device: str | None = None,
        *,
        use_flash_attn: bool = True,
        use_peri_ln: bool = False,
        rope_interleaved: bool = False,
        ff_cvmm_kernel: bool = False,
        ff_deepgemm_kernel: bool = False,
        ff_h100_fp8_train: bool = False,
        attn_cvmm_kernel: bool = False,
        init_std: float = 0.02,
        no_bias: bool = False,
        mup_config: MuPConfig | None = None,
    ) -> None:
        """Configure the DEPToEBlock.

        Parameters
        ----------
        d_model : int
            The dimensionality of the model.
        n_heads : int
            The number of attention heads.
        n_repeats : int
            The number of repeats for the Universal Transformer (UT) mechanism.
        n_blocks : int
            The number of blocks contained in the UT.
        use_peri_ln : bool
            If True, uses a peripheral layer normalization (PeriLN) strategy from
            https://arxiv.org/abs/2502.02732 instead of the layer normalization strategy
            proposed in https://arxiv.org/abs/2405.16039
        attn_heads_experts : int, optional
            The number of experts per head, by default 1.
        attn_heads_dropout : float, optional
            Dropout rate for attention, by default 0.0.
        attn_heads_projection_size : int, optional
            The size of the projection for each head, by default 64.
        attn_heads_experts_dropout : float, optional
            Dropout rate for experts in heads, by default 0.0.
        attn_heads_n_active_experts : int, optional
            The number of heads to use in the attention mechanism, by default 1.
        rope_base : int, optional
            Base frequency for the rotary embeddings, by default 10000.
        ff_n_experts : int, optional
            The number of experts in the feed-forward network, by default 1.
        ff_expert_size : int, optional
            The size of each expert in the feed-forward network, by default 64.
        ff_n_active_experts : int, optional
            The number of active experts in the feed-forward network, by default 1.
        ff_expert_dropout : float, optional
            Dropout rate for experts in the feed-forward network, by default 0.0.
        device : str | None, optional
            Device to place the model on (e.g., "cpu" or "cuda"), by default None.
        rope_interleaved : bool, optional
            If True, applies interleaved rotary positional encoding (GPT-J style)
            instead of the standard RoPE (GPT-NeoX style). Default is False.
        ff_cvmm_kernel : bool, optional
            If True, uses a custom kernel for the feed-forward network. Default is
            False.
        ff_deepgemm_kernel : bool, optional
            If True, uses the DeepGEMM kernel for the feed-forward network. Default is
            False.
        ff_h100_fp8_train : bool, optional
            If True, enables H100 FP8 training for the feed-forward network. Default is
            False.
        use_flash_attn : bool, optional
            If True, uses the Flash Attention implementation for the attention
            mechanism. This can significantly speed up the attention computation,
            especially for large models and long sequences. Default is True, meaning
            Flash Attention is used if available. If False, the standard PyTorch
            attention implementation is used.
        attn_cvmm_kernel : bool, optional
            If True, uses a custom kernel for the attention mechanism. Default is
            False.
        init_std : float
            The initial standard deviation for the weights in the block.
            This is used to initialize the weights of the block.
            Default is 0.02.
        no_bias : bool
            Whether to disable the biases in the model for linear layers/layernorms.
        mup_config : MuPConfig | None, optional
            Configuration for the MuP (Maximal Update Parametrisation) mechanism.
            If None, MuP is not enabled. Default is None.


        """
        self.d_model = d_model
        self.device = device

        self.n_repeats = n_repeats
        self.n_blocks = n_blocks

        self.n_heads = n_heads
        self.attention_dropout = attn_heads_dropout

        self.rope_interleaved = rope_interleaved
        self.rope_base = rope_base

        # Flag for using the peripheral layer normalization (PeriLN)
        self.use_peri_ln = use_peri_ln

        self.mup_config = MuPConfig() if mup_config is None else mup_config

        self.switch_head_config = SwitchHeadConfig(
            d_model=d_model,
            n_heads=n_heads,
            n_repeats=n_repeats,
            n_blocks=n_blocks,
            attn_heads_n_experts=attn_heads_experts,
            attn_heads_dropout=attn_heads_dropout,
            attn_heads_projection_size=attn_heads_projection_size,
            attn_heads_expert_dropout=attn_heads_experts_dropout,
            attn_heads_n_active_experts=attn_heads_n_active_experts,
            cvmm_kernel=attn_cvmm_kernel,
            use_flash_attn=use_flash_attn,
            use_peri_ln=use_peri_ln,
            init_std=init_std,
            no_bias=no_bias,
            mup_config=self.mup_config,
        )
        self.ff_n_experts = ff_n_experts
        self.ff_expert_size = ff_expert_size
        self.ff_n_active_experts = ff_n_active_experts
        self.ff_expert_dropout = ff_expert_dropout
        # TODO(anon): add activation and v_dim
        self.sigma_moe_config = SigmaMoEConfig(
            d_model=d_model,
            n_repeats=n_repeats,
            n_blocks=n_blocks,
            ff_n_experts=ff_n_experts,
            ff_expert_size=ff_expert_size,
            ff_n_active_experts=ff_n_active_experts,
            activation=torch.nn.functional.relu,
            v_dim=None,
            expert_dropout=ff_expert_dropout,
            cvmm_kernel=ff_cvmm_kernel,
            deepgemm_kernel=ff_deepgemm_kernel,
            h100_fp8_train=ff_h100_fp8_train,
            init_std=init_std,
            mup_config=self.mup_config,
        )


class DEPToEBlock(nn.Module):
    """A Transformer block with optional fused normalization and attention."""

    def __init__(
        self,
        config: DEPToEBlockConfig,
        rope_module: RotaryPositionalEncoding | None = None,
    ) -> None:
        """Initialize the DEPToEBlock.

        Parameters
        ----------
        config : DEPToEBlockConfig
            Configuration object containing parameters for the DEPToEBlock.
        rope_module : RotaryPositionalEncoding | None, optional
            A module for applying rotary positional embeddings, by default None.

        """
        super().__init__()
        self.config = config
        switch_head_config = config.switch_head_config
        sigma_moe_config = config.sigma_moe_config
        self.switch_head = SwitchHead(
            config=switch_head_config,
            rope_module=rope_module,
        )
        self.sigma_moe = SigmaMoE(config=sigma_moe_config)
        self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
        self.residual_scaling = (
            1.0
            / (
                self.config.mup_config.completep_depth_multiplier
                ** self.config.mup_config.completep_depth_alpha_exp
            )
            if self.config.mup_config.completep_depth_alpha_enabled
            else 1.0
        )

    def forward(
        self,
        x: torch.Tensor,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        attention_mask: torch.ByteTensor | None = None,
        *,
        output_attentions: bool = False,
    ) -> tuple[
        torch.Tensor,
        tuple[torch.Tensor, torch.Tensor] | None,
        torch.Tensor | None,
    ]:
        """Forward pass of the DEPToEBlock.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (batch_size, seq_len, d_model). This is the main data
            path.
        past_key_value : tuple[torch.Tensor, torch.Tensor] | None, optional
            Tuple containing past key and value tensors for caching, by default None.
        attention_mask : torch.ByteTensor | None, optional
            Attention mask tensor of shape (batch_size, seq_len), by default None.
        output_attentions : bool, optional
            Whether to return attention scores, by default False.

        Returns
        -------
        tuple[
            torch.Tensor,
            tuple[torch.Tensor, torch.Tensor] | None,
            torch.Tensor | None,
        ]
            A tuple containing:
            - The output tensor after attention and feed-forward network.
            - Attention scores tensor if applicable, otherwise None.
            - Updated past key and value tensors for caching, or None if not used.
            - Attention probabilities tensor if `output_attentions` is True, otherwise
                None.

        """
        attention_output, kv_cache, attn_probs = cast(
            "tuple[torch.Tensor, KVCache, torch.Tensor | None]",
            self.switch_head(
                x=x,
                mask=attention_mask,
                kv_cache=(
                    {
                        "v": past_key_value[1],
                        "k": past_key_value[0],
                    }
                    if past_key_value is not None
                    else None
                ),
                output_attentions=output_attentions,
            ),
        )
        # NOTE: In-place modification of `x` is not allowed by PyTorch Autograd
        x = x + self.attention_dropout(attention_output) * self.residual_scaling  # noqa: PLR6104
        up_and_down_projected_x = cast("torch.Tensor", self.sigma_moe(x))
        present_key_value = None if kv_cache is None else (kv_cache["k"], kv_cache["v"])
        return (
            x + up_and_down_projected_x * self.residual_scaling,
            present_key_value,
            attn_probs,
        )
