from transformers import PretrainedConfig


class FWKVConfig(PretrainedConfig):
    model_type = "fwkv"

    def __init__(
        self,
        vocab_size=50257,
        d_model=768,
        d_emb=384,
        n_layers=13,
        ffn_mult=4,
        head_size=64,
        wkv_floor=0.1,
        wkv_ceil=0.999,
        tau_min=2.0,
        tau_max=1000.0,
        seq_len=1024,
        ce_chunk=512,
        arch_version=2,
        tie_word_embeddings=True,
        **kwargs,
    ):
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.d_emb = d_emb
        self.n_layers = n_layers
        self.ffn_mult = ffn_mult
        self.head_size = head_size
        self.wkv_floor = wkv_floor
        self.wkv_ceil = wkv_ceil
        self.tau_min = tau_min
        self.tau_max = tau_max
        self.seq_len = seq_len
        self.ce_chunk = ce_chunk
        self.arch_version = arch_version

        # ─── Standard HF aliases ───
        self.num_hidden_layers = n_layers
        self.hidden_size = d_model
        self.num_attention_heads = 1       # dummy, not used by FWKV
        self.num_key_value_heads = 1       # dummy