from __future__ import annotations

from typing import Any

from transformers import PretrainedConfig


class HanseConfig(PretrainedConfig):
    model_type = "hanse"

    def __init__(
        self,
        vocab_size: int = 24_576,
        hidden_size: int = 640,
        num_layers: int = 14,
        layer_pattern: list[str] | tuple[str, ...] | None = None,
        num_query_heads: int = 10,
        num_kv_heads: int = 2,
        ffn_hidden_size: int = 1_792,
        max_seq_len: int = 2_048,
        rope_theta: float = 10_000.0,
        conv_kernel_size: int = 7,
        norm_eps: float = 1e-6,
        qk_norm: bool = True,
        tie_word_embeddings: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.layer_pattern = tuple(
            layer_pattern
            or ("A", "A", "C", "A", "A", "A", "C", "A", "A", "A", "C", "A", "A", "A")
        )
        self.num_query_heads = num_query_heads
        self.num_kv_heads = num_kv_heads
        self.ffn_hidden_size = ffn_hidden_size
        self.max_seq_len = max_seq_len
        self.rope_theta = rope_theta
        self.conv_kernel_size = conv_kernel_size
        self.norm_eps = norm_eps
        self.qk_norm = qk_norm

    @property
    def head_dim(self) -> int:
        return self.hidden_size // self.num_query_heads