"""CodVa-1 configuration for HuggingFace Transformers compatibility."""

from transformers import PretrainedConfig


class CodVa1Config(PretrainedConfig):
    model_type = "codva1"

    def __init__(
        self,
        vocab_size: int = 32064,
        d_model: int = 1536,
        n_layers: int = 28,
        n_heads: int = 24,
        n_kv_heads: int = 6,
        max_len: int = 2048,
        rope_theta: float = 5_000_000.0,
        ffn_hidden: int = 4096,
        use_moe: bool = True,
        moe_experts: int = 16,
        moe_top_k: int = 2,
        moe_shared: int = 2,
        moe_hidden: int = 1024,
        moe_every: int = 2,
        use_qk_norm: bool = True,
        use_structural_bias: bool = True,
        n_struct_rel: int = 4,
        rms_norm_eps: float = 1e-6,
        tie_word_embeddings: bool = True,
        fim_pre_id: int = -1,
        fim_suf_id: int = -1,
        fim_mid_id: int = -1,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.max_len = max_len
        self.rope_theta = rope_theta
        self.ffn_hidden = ffn_hidden

        self.use_moe = use_moe
        self.moe_experts = moe_experts
        self.moe_top_k = moe_top_k
        self.moe_shared = moe_shared
        self.moe_hidden = moe_hidden
        self.moe_every = moe_every

        self.use_qk_norm = use_qk_norm
        self.use_structural_bias = use_structural_bias
        self.n_struct_rel = n_struct_rel
        self.rms_norm_eps = rms_norm_eps

        # FIM special token ids (padded into vocab during training)
        self.fim_pre_id = fim_pre_id
        self.fim_suf_id = fim_suf_id
        self.fim_mid_id = fim_mid_id

        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

    # Newer `transformers` internals (cache utils, generation config, etc.)
    # look for these standard attribute names regardless of custom naming,
    # even when the model declares no cache support. Alias them through.
    @property
    def num_hidden_layers(self):
        return self.n_layers

    @property
    def num_attention_heads(self):
        return self.n_heads

    @property
    def hidden_size(self):
        return self.d_model