"""MPTCompletePBlock class, which is a specialized.

block for the MPT architecture that incorporates CompleteP (Complete Parametrisation)
mechanisms. It scales the residual branches to control the contribution of each block
to the overall model output, allowing for more flexible and efficient training of large
language models.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from llmfoundry.models.layers.blocks import MPTBlock
from llmfoundry.models.layers.layer_builders import build_norm

if TYPE_CHECKING:
    import torch

log = logging.getLogger(__name__)


class MPTCompletePBlock(MPTBlock):
    """MPTBlock that scales residual branches for CompleteP."""

    def __init__(  # noqa: PLR0913
        self,
        depth_multiplier: float = 1.0,
        depth_alpha_exp: float = 1.0,
        norm_type: str = "low_precision_layernorm",
        norm_eps: float = 1e-05,
        *,
        use_peri_norm: bool = False,
        depth_alpha_enabled: bool = False,
        device: str | None = None,
        **kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize the MPTCompletePBlock.

        Parameters
        ----------
        depth_multiplier : float
            The multiplier for the depth of the block, controlling how much each block
            contributes to the overall model output.
        depth_alpha_exp : float
            The exponent for scaling the block's contribution, allowing for flexible
            adjustment of the residual scaling.
        norm_type : str
            The type of normalization to apply, e.g., "layernorm" or "low_precision
            layernorm".
        norm_eps : float
            The epsilon value for normalization, preventing division by zero.
        use_peri_norm : bool
            Whether to use peri normalization, which can help stabilize training
        depth_alpha_enabled : bool
            Whether to enable depth alpha scaling, which adjusts the contribution of
            each block based on its depth in the model.
        device : str | None
            The device on which to place the block's parameters, e.g., "cpu" or "cuda".
        **kwargs : Any
            Additional keyword arguments for the parent class initialization.

        """
        super().__init__(
            norm_type=norm_type,
            norm_eps=norm_eps,
            device=device,
            **kwargs,
        )
        self.residual_scaling = (
            1.0 / (depth_multiplier**depth_alpha_exp) if depth_alpha_enabled else 1.0
        )
        self.post_attn_norm = None
        self.post_ffn_norm = None
        if use_peri_norm:
            self.post_attn_norm = build_norm(
                name=norm_type.lower(),
                normalized_shape=kwargs["d_model"],
                eps=norm_eps,
                device=device,
            )
            self.post_ffn_norm = build_norm(
                name=norm_type.lower(),
                normalized_shape=kwargs["d_model"],
                eps=norm_eps,
                device=device,
            )

    def forward(  # noqa: PLR0913, PLR0917
        self,
        x: torch.Tensor,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        attn_bias: torch.Tensor | None = None,
        rotary_emb_w_meta_info: dict | None = None,
        attention_mask: torch.ByteTensor | None = None,
        is_causal: bool = True,  # noqa: FBT001, FBT002
        output_attentions: bool = False,  # noqa: FBT001, FBT002
        alibi_slopes: torch.Tensor | None = None,
        flash_attn_padding_info: dict[str, torch.Tensor] | None = None,
        prev_layer_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
        key_value_states: torch.Tensor | None = None,
        pos_id_within_seq: torch.Tensor | None = None,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor | None,
        tuple[torch.Tensor, torch.Tensor] | None,
    ]:
        """Forward pass of the MPTCompletePBlock.

        Parameters
        ----------
        x : torch.Tensor
            The input tensor to the block, typically the output from the previous layer.
        past_key_value : tuple[torch.Tensor, torch.Tensor] | None
            The cached key and value tensors from previous attention layers, used for
            efficient attention computation.
        attn_bias : torch.Tensor | None
            The attention bias tensor, which can be used to modify the attention scores.
        rotary_emb_w_meta_info : dict | None
            Metadata for rotary embeddings, if used in the attention mechanism.
        attention_mask : torch.ByteTensor | None
            The attention mask to apply, typically used to prevent attention to padding
            tokens or future tokens in causal attention.
        is_causal : bool, optional
            Whether the attention is causal (default is True). If True, prevents
            attending to future tokens.
        output_attentions : bool, optional
            Whether to return attention weights (default is False). If True, the method
            will return the attention weights along with the output tensor.
        alibi_slopes : torch.Tensor | None
            The slopes for ALiBi (Attention with Linear Biases) attention, if used.
        flash_attn_padding_info : dict[str, torch.Tensor] | None
            Padding information for Flash Attention, if applicable.
        prev_layer_key_value : tuple[torch.Tensor, torch.Tensor] | None
            The key and value tensors from the previous layer, if applicable.
        key_value_states : torch.Tensor | None
            The key and value states for the attention mechanism, if applicable.
        pos_id_within_seq : torch.Tensor | None
            Positional IDs within the sequence, used for rotary embeddings or other
            positional encodings.

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor, torch.Tensor]
                | None]
            A tuple containing:
            - The output tensor after applying the attention and feed-forward layers.
            - The attention weights if `output_attentions` is True, otherwise None.
            - The updated past key-value tuple, which can be used for subsequent
                attention computations.

        """
        extra_kwargs = {}
        if prev_layer_key_value is not None:
            extra_kwargs["prev_layer_key_value"] = prev_layer_key_value
        if key_value_states is not None:
            extra_kwargs["key_value_states"] = key_value_states
        if pos_id_within_seq is not None:
            extra_kwargs["pos_id_within_seq"] = pos_id_within_seq

        if self.fuse_norm_attn_norm:
            a = self.norm_attn_norm.norm_1(x)
            b, attn_weights, past_key_value = self.norm_attn_norm.attn(
                a,
                past_key_value=past_key_value,
                attn_bias=attn_bias,
                rotary_emb_w_meta_info=rotary_emb_w_meta_info,
                attention_mask=attention_mask,
                is_causal=is_causal,
                needs_weights=output_attentions,
                alibi_slopes=alibi_slopes,
                flash_attn_padding_info=flash_attn_padding_info,
                **extra_kwargs,
            )
            b = self.post_attn_norm(b) if self.post_attn_norm is not None else b
            x = (  # noqa: PLR6104
                x
                + self.norm_attn_norm.resid_attn_dropout(
                    b,
                )
                * self.residual_scaling
            )
            m = x
            if self.norm_attn_norm.norm_2 is not None:
                m = self.norm_attn_norm.norm_2(x)
        else:
            a = self.norm_1(x)
            b, attn_weights, past_key_value = self.attn(
                a,
                past_key_value=past_key_value,
                attn_bias=attn_bias,
                rotary_emb_w_meta_info=rotary_emb_w_meta_info,
                attention_mask=attention_mask,
                is_causal=is_causal,
                needs_weights=output_attentions,
                alibi_slopes=alibi_slopes,
                flash_attn_padding_info=flash_attn_padding_info,
                **extra_kwargs,
            )
            b = self.post_attn_norm(b) if self.post_attn_norm is not None else b
            x = x + self.resid_attn_dropout(b) * self.residual_scaling  # noqa: PLR6104
            m = x
            if self.norm_2 is not None:
                m = self.norm_2(x)

        n = self.apply_ffn(attention_mask, m)
        n = self.post_ffn_norm(n) if self.post_ffn_norm is not None else n
        x = (
            x.to(device=n.device)
            + self.resid_ffn_dropout(n).to(
                device=n.device,
            )
            * self.residual_scaling
        )
        return x, attn_weights, past_key_value
