"""Module implements the configuration for the DEPToE model."""

import copy
from dataclasses import asdict
from typing import Any

from transformers.configuration_utils import PretrainedConfig

from mt_dao.models.deptoe.deptoe_block import DEPToEBlockConfig
from mt_dao.models.deptoe.mup_config import MuPConfig

INIT_CONFIG_DEFAULTS: dict = {
    "name": "kaiming_normal_",
    "fan_mode": "fan_in",
    "init_nonlinearity": "relu",
    "init_div_is_residual": True,
    "emb_init_std": None,
    "emb_init_uniform_lim": None,
    "init_std": None,
    "init_gain": 0.0,
}


class DEPToEConfig(PretrainedConfig):
    """The configuration class for the DEPToE model.

    It inherits from `PretrainedConfig` and is used to define the model's architecture,
    hyperparameters, and other settings. This class allows for easy configuration and
    instantiation of the DEPToE model with various parameters such as the number of
    layers, number of attention heads, model dimension, expansion ratio, and more. Its
    inheritance from `PretrainedConfig` enables compatibility with the Hugging Face
    Transformers library, allowing for seamless integration with pre-trained models and
    tokenizers.

    """

    model_type = "deptoe"

    def __init__(  # noqa: PLR0913, PLR0917
        self,
        d_model: int = 2048,
        n_heads: int = 16,
        n_layers: int = 24,
        max_seq_len: int = 2048,
        vocab_size: int = 50368,
        embedding_dropout: float = 0.0,
        init_device: str = "cpu",
        logit_scale: float | str | None = None,
        embedding_fraction: float = 1.0,
        out_norm_type: str = "low_precision_layernorm",
        out_norm_eps: float = 1e-05,
        init_config: dict | None = None,
        final_logit_softcapping: float | None = None,
        # SwitchHeadConfig parameters
        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,
        # SigmaMoE parameters
        ff_n_experts: int = 1,
        ff_expert_size: int = 64,
        ff_n_active_experts: int = 1,
        # MoEUT regularization parameters
        entropy_reg: float = 0.01,
        att_entropy_reg: float = 0.001,
        # UT number of repeats
        n_repeats: int = 1,
        # RoPE parameters
        rope_base: int = 10000,
        *,
        use_flash_attn: bool = True,
        attn_xpos_scale_base: int | None = None,
        use_torch_apply_rotary_emb: bool = False,
        use_peri_ln: bool = False,
        use_embedding_peri_ln: bool = False,
        attn_cvmm_kernel: bool = False,
        ff_cvmm_kernel: bool = False,
        rope_interleaved: bool = False,
        use_cache: bool = False,
        tie_word_embeddings: bool = True,
        no_bias: bool = True,
        # MuP parameters,
        mup_config: dict[str, Any] | None = None,
        **kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize the DEPToE configuration class.

        Parameters
        ----------
        d_model : int
            The size of the embedding dimension of the model. The default is 2048,
            which is a common choice for many medium-sized transformer models.
        n_heads : int
            The number of attention heads. The default is 16,
            which is a common choice for many medium-sized transformer models.
        n_layers : int
            The number of layers/blocks in the model. The default is 24,
            which is a common choice for many medium-sized transformer models.
        max_seq_len : int
            The maximum sequence length that the model supports. The default is 2048,
            which is a common choice for many medium-sized transformer models.
        vocab_size : int
            The size of the vocabulary for the model. The default is 50368,
            which is a common size for many language models.
        embedding_dropout : float
            The dropout probability for the embedding layer. The default is 0.0,
            meaning no dropout is applied to the embedding layer.
        init_device : str
            The device to use for parameter initialization (e.g., "cpu" or "cuda").
            The default is "cpu", meaning parameters are initialized on the CPU.
        logit_scale : float | str | None
            If not None, scale the logits by this value. If "inv_sqrt_d_model", scale by
            the inverse square root of `d_model`. If None, no scaling is applied. The
            default is None, meaning no scaling applied.
        embedding_fraction : float
            The fraction to scale the gradients of the embedding layer by. This was
            first proposed on page 7 of the GLM-130B paper https://arxiv.org/abs/2210.02414.
            A value of 1.0 means no scaling. Default is 1.0, meaning no scaling applied.
        out_norm_type : str
            The type of normalization to use for the output layer. Options include
            "layernorm", "low_precision_layernorm", "rmsnorm", or "none". If "none",
            no normalization is applied to the output layer. Default is
            "low_precision_layernorm".
        out_norm_eps : float
            The epsilon value for the normalization layer. This is used to avoid
            division by zero in the normalization calculations.
        init_config : dict | None
            A dictionary used to configure the model initialization:
                init_config.name (str): The parameter initialization scheme to use.
                    Options are: 'default_', 'baseline_', 'kaiming_uniform_',
                    'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_',
                    or 'xavier_normal_'. These mimic the parameter initialization
                    methods in PyTorch.
                init_div_is_residual (int | float | str | bool): Value to divide initial
                    weights by if ``module._is_residual`` is True.
                emb_init_std (float | None): The standard deviation of the normal
                    distribution used to initialize the embedding layer.
                emb_init_uniform_lim (Tuple[float, float] | float | None): The lower and
                    upper limits of the uniform distribution used to initialize the
                    embedding layer. Mutually exclusive with ``emb_init_std``.
                init_std (float): The standard deviation of the normal distribution used
                    to initialize the model, if using the 'baseline_' parameter
                    initialization scheme.
                init_gain (float): The gain to use for parameter initialization with
                    'kaiming_normal_' or 'xavier_normal_' initialization schemes.
                fan_mode (str): The fan mode to use for parameter initialization with
                    kaiming initialization schemes.
                init_nonlinearity (str): The nonlinearity to use for parameter
                    initialization with kaiming initialization schemes.
            ---
            See llmfoundry.models.utils.param_init_fns.py for info on other param init
            config options
        final_logit_softcapping : float | None
            Softcapping threshold for final logit. Set to None to disable
            (default value None). Please see https://arxiv.org/pdf/2403.08295 for more
            details.
        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
        use_embedding_peri_ln : bool
            If True, applies a peripheral layer normalization (PeriLN) after the token
            embedding as optionally proposed in https://arxiv.org/abs/2405.16039.
        attn_heads_experts : int
            The number of experts in the SwitchHead. Default is 1, meaning no
            experts are used.
        attn_heads_dropout : float
            The dropout probability for the query vectors in the attention layer.
            Default is 0.0, meaning no dropout is applied to the attention layer.
        attn_heads_projection_size : int
            The size of the projection for each attention head. Default is 64, which is
            a common choice for many transformer models.
        attn_heads_experts_dropout : float
            The dropout probability for the experts scores in the SwitchHead.
            Default is 0.0, meaning no dropout is applied to the experts.
        attn_heads_n_active_experts : int
            The number of attention heads to use in the SwitchHead. Default is 1,
            meaning a single attention head is used.
        use_flash_attn : bool
            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_xpos_scale_base: int | None
            The base scale for the XPOS (eXtra Positioning) encoding in the attention
            mechanism. This is used to control the scale of the positional encodings.
            Default is 512, which is a common choice for many transformer models.
        use_torch_apply_rotary_emb : bool
            If True, uses the PyTorch implementation of apply_rotary_emb instead of the
            Flash Attention implementation.
        ff_n_experts : int
            The number of experts in the feed-forward network. Default is 1, meaning no
            experts are used in the feed-forward network.
        ff_expert_size : int
            The size of each expert in the feed-forward network. Default is 64, which is
            a common choice for many transformer models.
        ff_n_active_experts : int
            The number of active experts in the feed-forward network. This is used to
            control how many experts are used during training and inference. Default is
            1, meaning only one expert is active at a time.
        entropy_reg : float
            The regularization coefficient for the entropy loss in the SigmaMoE layer.
            This is used to encourage diversity among the experts in the SigmaMoE layer.
            Default is 0.01, which is a common choice for many transformer models.
        att_entropy_reg : float
            The regularization coefficient for the entropy loss in the SwitchHead
            layer. This is used to encourage diversity among the attention heads in the
            SwitchHead layer. Default is 0.001, which is a common choice for many
            transformer models.
        n_repeats : int
            The number of repeats for the UT (Universal Transformer) mechanism. This
            controls how many times the input sequence is processed in the UT mechanism.
            Default is 1, meaning no repeats are used.
        rope_base : int
            The base frequency for the rotary positional encoding. This is used to
            control the scale of the positional encodings. Default is 10000, which is a
            common choice in transformer models.
        rope_interleaved : bool
            If True, applies interleaved rotary positional encoding (GPT-J style)
            instead of the standard RoPE (GPT-NeoX style). This is useful for handling
            long sequences and improving the model's ability to capture positional
            information. Default is False, meaning the standard RoPE is used.
        attn_cvmm_kernel : bool
            If True, uses a custom kernel for the attention mechanism. This can improve
            performance for large models by optimizing the forward pass. Default is
            False, meaning the standard PyTorch implementation is used.
        ff_cvmm_kernel : bool
            If True, uses a custom kernel for the feed-forward network. This can improve
            performance for large models by optimizing the forward pass. Default is
            False, meaning the standard PyTorch implementation is used.
        use_cache : bool
            Whether to use caching for the attention mechanism. Default is False,
            meaning caching is not used.
        tie_word_embeddings : bool
            Whether to tie the input embedding and output layers.
        no_bias : bool
            Whether to disable the biases in the model for linear layers/layernorms.
        mup_config : dict | None
            Configuration for MuP (Maximal Update Parametrisation) and CompleteP.
            If None, defaults to a standard configuration with `mup_enabled` False.
            Structure:
                mup_enabled : bool = False
                    Whether to enable MuP.
                mup_disable_attention_scaling : bool = False
                    Whether to disable attention scaling in MuP.
                mup_disable_hidden_lr_scaling: bool = False
                    Whether to disable hidden layer learning rate scaling in MuP.
                mup_width_multiplier : float = 1.0
                    The width multiplier for MuP, controlling hyperparameter scaling.
                    Default is 1.0, meaning we are training the base size.
                mup_input_alpha : float = 1.0
                    The input alpha for MuP, controlling input layer output scaling.
                mup_output_alpha : float = 1.0
                    The output alpha for MuP, controlling output layer scaling.
                mup_scale_expert_sel: bool = False
                    Whether to scale expert selection layer in the MuP way.
                    Defaults to False since the selection does not impact
                    the scale of the outputs for the transformer blocks.
                completep_depth_alpha_enabled : bool = True
                    Whether to enable depth alpha in CompleteP.
                    Meaning the output of the blocks will be scaled
                    based on the depth multiplier and exponent
                    before being merged to the residual path.
                completep_depth_multiplier : float = 1.0
                    The depth multiplier for CompleteP, scaling a block's contribution.
                    Default means no scaling since we are training the base size.
                completep_depth_alpha_exp : float = 1.0
                    The exponent for scaling a block's contribution in CompleteP.
                    Scaling is by 1/(depth_multiplier ** depth_alpha_exp).
                completep_eps_scaling_enabled : bool = True
                    Whether to enable epsilon scaling for ADAM/ADOPT in CompleteP.
                    If used epsilon is scaled based on both width and depth.
        kwargs : Any
            Other relevant keyword arguments.

        Raises
        ------
            ValueError: If `max_seq_len` is not an integer.

        """
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        if max_seq_len != int(max_seq_len):
            msl_type = type(max_seq_len)
            msg = f"max_seq_len={max_seq_len} must be an int, not a {msl_type}."
            raise ValueError(msg)
        self.max_seq_len = int(max_seq_len)
        self.vocab_size = vocab_size
        self.embedding_dropout = embedding_dropout

        self.init_device = init_device
        self.logit_scale = logit_scale
        self.embedding_fraction = embedding_fraction
        self.out_norm_type = out_norm_type
        self.out_norm_eps = out_norm_eps
        self.use_cache = use_cache
        self.init_config = (
            init_config
            if init_config is not None
            else copy.deepcopy(
                INIT_CONFIG_DEFAULTS,
            )
        )

        self.final_logit_softcapping = final_logit_softcapping

        self.no_bias = no_bias

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

        # SwitchHeadConfig parameters
        self.attn_heads_experts = attn_heads_experts
        self.attn_heads_dropout = attn_heads_dropout
        self.attn_heads_projection_size = attn_heads_projection_size
        self.attn_heads_experts_dropout = attn_heads_experts_dropout
        self.attn_heads_n_active_experts = attn_heads_n_active_experts
        self.attn_cvmm_kernel = attn_cvmm_kernel

        self.use_flash_attn = use_flash_attn
        self.attn_xpos_scale_base = attn_xpos_scale_base
        self.use_torch_apply_rotary_emb = use_torch_apply_rotary_emb

        # RoPE parameters
        self.rope_interleaved = rope_interleaved
        self.rope_base = rope_base

        # SigmaMoE parameters
        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_cvmm_kernel = ff_cvmm_kernel

        # MoEUT regularization parameters
        self.entropy_reg = entropy_reg
        self.att_entropy_reg = att_entropy_reg

        # UT number of repeats
        self.n_repeats = n_repeats

        mup_cfg_dataclass = (
            MuPConfig() if mup_config is None else MuPConfig(**mup_config)
        )

        self.mup_config: dict[str, Any] = asdict(mup_cfg_dataclass)

        if (
            mup_cfg_dataclass.mup_enabled
            and mup_cfg_dataclass.completep_depth_alpha_enabled
            and self.n_repeats > 1
        ):
            msg = (
                "MuP and CompleteP depth alpha are not compatible with n_repeats > 1. "
                "Please set n_repeats to 1 or disable CompleteP depth alpha."
            )
            raise ValueError(msg)

        kwargs.pop("name", None)
        kwargs.pop("loss_fn", None)
        # NOTE: tie_word_embeddings is set in Huggingface's PretrainedConfig __init__
        super().__init__(
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

        self.validate_config()

        # Create the DEPToEBlockConfig object
        self.deptoe_block_config: DEPToEBlockConfig = DEPToEBlockConfig(
            d_model=self.d_model,
            n_repeats=self.n_repeats,
            n_blocks=self.n_layers,
            use_peri_ln=self.use_peri_ln,
            # SwitchHeadConfig parameters
            n_heads=self.n_heads,
            attn_heads_experts=attn_heads_experts,
            attn_heads_dropout=attn_heads_dropout,
            attn_heads_projection_size=attn_heads_projection_size,
            attn_heads_experts_dropout=attn_heads_experts_dropout,
            attn_heads_n_active_experts=attn_heads_n_active_experts,
            attn_cvmm_kernel=attn_cvmm_kernel,
            use_flash_attn=use_flash_attn,
            # RoPE arguments
            rope_interleaved=rope_interleaved,
            rope_base=rope_base,
            # Sigma MoE arguments
            ff_n_experts=ff_n_experts,
            ff_expert_size=ff_expert_size,
            ff_n_active_experts=ff_n_active_experts,
            ff_cvmm_kernel=ff_cvmm_kernel,
            device=self.init_device,
            init_std=self.init_config.get("init_std", 0.02),
            no_bias=no_bias,
            # MuP arguments
            mup_config=mup_cfg_dataclass,
        )

    def _set_config_defaults(
        self,
        config: dict[str, Any],
        config_defaults: dict[str, Any],
    ) -> dict[str, Any]:
        # set config defaults
        for k, v in config_defaults.items():
            if k not in config:
                config[k] = v
            elif isinstance(v, dict):
                # recursively set default values for any sub-dicts
                config[k] = self._set_config_defaults(
                    config[k] if (config[k] is not None) else {},
                    v,
                )
        return config

    def validate_config(self) -> None:
        """Validate the configuration parameters.

        Raises
        ------
        ValueError : If any of the configuration parameters are invalid, such as:
            - `d_model` not being divisible by `n_heads`
            - `embedding_dropout` not being between 0 and 1
            - `embedding_fraction` not being between 0 (exclusive) and 1 (inclusive)
            - `logit_scale` not being a recognized option
            - `init_config` not having a 'name' key set.

        """
        self.init_config = self._set_config_defaults(
            self.init_config,
            INIT_CONFIG_DEFAULTS,
        )

        if self.d_model % self.n_heads != 0:
            msg = f"d_model={self.d_model} must be divisible by n_heads={self.n_heads}"
            raise ValueError(msg)
        if self.embedding_dropout < 0 or self.embedding_dropout > 1:
            msg = (
                f"embedding_dropout={self.embedding_dropout} is a probability and"
                " must be between 0 and 1"
            )
            raise ValueError(msg)
        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
            msg = (
                f"model.embedding_fraction={self.embedding_fraction} must be"
                " between 0 (exclusive) and 1 (inclusive)!"
            )
            raise ValueError(msg)
        if (
            isinstance(
                self.logit_scale,
                str,
            )
            and self.logit_scale != "inv_sqrt_d_model"
        ):
            msg = (
                f"{self.logit_scale=} is not recognized as an option; use numeric"
                " value or 'inv_sqrt_d_model'."
            )
            raise ValueError(msg)
        if self.init_config.get("name", None) is None:
            msg = f"{self.init_config=} 'name' needs to be set."
            raise ValueError(msg)
