from __future__ import annotations

from typing import Any

from transformers import PretrainedConfig


class LizzyConfig(PretrainedConfig):
    model_type = "lizzy"
    keys_to_ignore_at_inference = ["past_key_values"]
    base_model_tp_plan = {
        "layers.*.self_attn.q_proj": "colwise",
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
        "lm_head": "colwise",
    }
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

    def __init__(
        self,
        vocab_size: int = 32000,
        hidden_size: int = 4096,
        intermediate_size: int = 11008,
        num_hidden_layers: int = 32,
        num_attention_heads: int = 32,
        num_key_value_heads: int | None = None,
        max_position_embeddings: int = 2048,
        head_dim: int | None = None,
        hidden_act: str = "silu",
        norm_type: str = "rmsnorm",
        norm_eps: float = 1e-6,
        norm_has_bias: bool = False,
        use_pre_attn_norm: bool = True,
        use_pre_mlp_norm: bool = True,
        use_post_attn_norm: bool = False,
        use_post_mlp_norm: bool = False,
        mlp_type: str = "gated",
        attention_bias: bool = False,
        mlp_bias: bool = False,
        position_embedding_type: str = "rope",
        rope_theta: float = 10000.0,
        rope_scaling: dict[str, Any] | None = None,
        rope_layer_flags: list[bool] | None = None,
        no_rope_layer_interval: int | None = None,
        rope_type_overrides: dict[str, str] | None = None,
        layer_types: list[str] | None = None,
        layer_layouts: list[str] | None = None,
        sliding_window: int | None = None,
        linear_num_key_heads: int | None = None,
        linear_num_value_heads: int | None = None,
        linear_key_head_dim: int | None = None,
        linear_value_head_dim: int | None = None,
        linear_a_log_min: float | None = None,
        linear_a_log_max: float | None = None,
        linear_dt_min: float | None = None,
        linear_dt_max: float | None = None,
        linear_dt_init_floor: float | None = None,
        linear_conv_kernel_dim: int | None = None,
        linear_allow_neg_eigval: bool | None = None,
        use_qk_norm: bool = False,
        qk_norm_type: str = "rmsnorm",
        attention_dropout: float = 0.0,
        resid_dropout: float = 0.0,
        embd_dropout: float = 0.0,
        initializer_range: float = 0.02,
        bos_token_id: int | None = None,
        eos_token_id: int | None = None,
        pad_token_id: int | None = None,
        use_cache: bool = True,
        tie_word_embeddings: bool = False,
        **kwargs,
    ) -> None:
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads
        if head_dim is None:
            head_dim = hidden_size // num_attention_heads
        if no_rope_layer_interval is not None:
            no_rope_layer_interval = int(no_rope_layer_interval)
            if no_rope_layer_interval <= 0:
                no_rope_layer_interval = None
        if layer_types is None:
            layer_types = ["full_attention"] * int(num_hidden_layers)
        if layer_layouts is None:
            if use_post_attn_norm or use_post_mlp_norm:
                layer_layouts = ["decoder_postnorm"] * int(num_hidden_layers)
            else:
                layer_layouts = ["decoder_prenorm"] * int(num_hidden_layers)
        if rope_layer_flags is None:
            rope_enabled = position_embedding_type == "rope"
            if rope_enabled and no_rope_layer_interval is not None:
                rope_layer_flags = [
                    ((layer_idx + 1) % no_rope_layer_interval) != 0
                    for layer_idx in range(int(num_hidden_layers))
                ]
            else:
                rope_layer_flags = [rope_enabled] * int(num_hidden_layers)

        normalized_rope_scaling = None
        if rope_scaling is not None:
            normalized_rope_scaling = dict(rope_scaling)
            for field_name in (
                "factor",
                "attention_factor",
                "beta_fast",
                "beta_slow",
            ):
                if normalized_rope_scaling.get(field_name) is not None:
                    normalized_rope_scaling[field_name] = float(
                        normalized_rope_scaling[field_name]
                    )
            if (
                normalized_rope_scaling.get("original_max_position_embeddings")
                is not None
            ):
                normalized_rope_scaling["original_max_position_embeddings"] = int(
                    normalized_rope_scaling["original_max_position_embeddings"]
                )

        # Transformers validates RoPE settings during PretrainedConfig
        # initialization, so publish the rope-critical fields before
        # calling `super().__init__()`.
        self.max_position_embeddings = int(max_position_embeddings)
        self.rope_theta = float(rope_theta)
        self.rope_scaling = normalized_rope_scaling

        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
        self.vocab_size = int(vocab_size)
        self.hidden_size = int(hidden_size)
        self.intermediate_size = int(intermediate_size)
        self.num_hidden_layers = int(num_hidden_layers)
        self.num_attention_heads = int(num_attention_heads)
        self.num_key_value_heads = int(num_key_value_heads)
        self.max_position_embeddings = int(max_position_embeddings)
        self.head_dim = int(head_dim)
        self.hidden_act = str(hidden_act)
        self.norm_type = str(norm_type)
        self.norm_eps = float(norm_eps)
        self.norm_has_bias = bool(norm_has_bias)
        self.use_pre_attn_norm = bool(use_pre_attn_norm)
        self.use_pre_mlp_norm = bool(use_pre_mlp_norm)
        self.use_post_attn_norm = bool(use_post_attn_norm)
        self.use_post_mlp_norm = bool(use_post_mlp_norm)
        self.mlp_type = str(mlp_type)
        self.attention_bias = bool(attention_bias)
        self.mlp_bias = bool(mlp_bias)
        self.position_embedding_type = str(position_embedding_type)
        self.rope_theta = float(rope_theta)
        self.rope_scaling = normalized_rope_scaling
        self.no_rope_layer_interval = no_rope_layer_interval
        self.rope_type_overrides = {
            str(key): str(value)
            for key, value in dict(rope_type_overrides or {}).items()
        }
        self.layer_types = list(layer_types)
        self.layer_layouts = [str(item) for item in layer_layouts]
        self.rope_layer_flags = [bool(item) for item in rope_layer_flags]
        self.sliding_window = sliding_window
        self.linear_num_key_heads = (
            None
            if linear_num_key_heads is None
            else int(linear_num_key_heads)
        )
        self.linear_num_value_heads = (
            None
            if linear_num_value_heads is None
            else int(linear_num_value_heads)
        )
        self.linear_key_head_dim = (
            None
            if linear_key_head_dim is None
            else int(linear_key_head_dim)
        )
        self.linear_value_head_dim = (
            None
            if linear_value_head_dim is None
            else int(linear_value_head_dim)
        )
        self.linear_a_log_min = (
            None if linear_a_log_min is None else float(linear_a_log_min)
        )
        self.linear_a_log_max = (
            None if linear_a_log_max is None else float(linear_a_log_max)
        )
        self.linear_dt_min = (
            None if linear_dt_min is None else float(linear_dt_min)
        )
        self.linear_dt_max = (
            None if linear_dt_max is None else float(linear_dt_max)
        )
        self.linear_dt_init_floor = (
            None
            if linear_dt_init_floor is None
            else float(linear_dt_init_floor)
        )
        self.linear_conv_kernel_dim = (
            None
            if linear_conv_kernel_dim is None
            else int(linear_conv_kernel_dim)
        )
        self.linear_allow_neg_eigval = (
            None
            if linear_allow_neg_eigval is None
            else bool(linear_allow_neg_eigval)
        )
        self.use_qk_norm = bool(use_qk_norm)
        self.qk_norm_type = str(qk_norm_type)
        self.attention_dropout = float(attention_dropout)
        self.resid_dropout = float(resid_dropout)
        self.embd_dropout = float(embd_dropout)
        self.initializer_range = float(initializer_range)
        self.use_cache = bool(use_cache)
        self.rms_norm_eps = self.norm_eps
        self.dtype = None
