"""RavenGuard HF config (trust_remote_code).

Mirrors the Netis Amniota ``GPTConfig`` fields 1:1 so ``RavenGuardForCausalLM`` can
rebuild the exact Netis Amniota ``GPT`` used to train the checkpoint. All architecture
knobs are stored verbatim from the checkpoint ``meta_*.json`` ``model_config``.
"""
from transformers import PretrainedConfig


class RavenGuardConfig(PretrainedConfig):
    model_type = "ravenguard"

    # names of the fields that map straight onto netis_amniota.model_config.GPTConfig
    gpt_config_fields = [
        "sequence_len", "vocab_size", "n_layer", "n_head", "n_kv_head", "n_embd",
        "window_pattern", "msa_block_size", "msa_top_k_blocks", "msa_local_window",
        "msa_sink_blocks", "msa_index_mode", "msa_kernel", "sparse_token_budget",
        "sparse_block_size", "sparse_block_budget", "sparse_index_mode",
        "sparse_index_heads", "sparse_index_dim", "ve_n_unique", "rope_base",
        "doc_mask", "bos_token_id", "grad_checkpoint", "n_experts",
        "n_experts_active", "n_shared_experts", "moe_ffn_mult", "moe_aux_coef",
        "moe_zloss_coef", "moe_norm_topk", "emo_doc_pool", "emo_bos_id",
    ]

    def __init__(
        self,
        sequence_len=2048,
        vocab_size=65536,
        n_layer=20,
        n_head=10,
        n_kv_head=10,
        n_embd=1280,
        window_pattern="SSSL",
        msa_block_size=64,
        msa_top_k_blocks=16,
        msa_local_window=128,
        msa_sink_blocks=1,
        msa_index_mode="pooled_k",
        msa_kernel="sdpa",
        sparse_token_budget=512,
        sparse_block_size=64,
        sparse_block_budget=16,
        sparse_index_mode="proxy_qk",
        sparse_index_heads=1,
        sparse_index_dim=64,
        ve_n_unique=3,
        rope_base=100000,
        doc_mask=False,
        bos_token_id=-1,
        grad_checkpoint=False,
        n_experts=0,
        n_experts_active=0,
        n_shared_experts=0,
        moe_ffn_mult=2.0,
        moe_aux_coef=0.01,
        moe_zloss_coef=0.001,
        moe_norm_topk=True,
        emo_doc_pool=False,
        emo_bos_id=-1,
        tie_word_embeddings=False,
        **kwargs,
    ):
        # let HF set its reserved attrs first, then our fields win
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
        self.sequence_len = sequence_len
        self.vocab_size = vocab_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_kv_head = n_kv_head
        self.n_embd = n_embd
        self.window_pattern = window_pattern
        self.msa_block_size = msa_block_size
        self.msa_top_k_blocks = msa_top_k_blocks
        self.msa_local_window = msa_local_window
        self.msa_sink_blocks = msa_sink_blocks
        self.msa_index_mode = msa_index_mode
        self.msa_kernel = msa_kernel
        self.sparse_token_budget = sparse_token_budget
        self.sparse_block_size = sparse_block_size
        self.sparse_block_budget = sparse_block_budget
        self.sparse_index_mode = sparse_index_mode
        self.sparse_index_heads = sparse_index_heads
        self.sparse_index_dim = sparse_index_dim
        self.ve_n_unique = ve_n_unique
        self.rope_base = rope_base
        self.doc_mask = doc_mask
        self.bos_token_id = bos_token_id
        self.grad_checkpoint = grad_checkpoint
        self.n_experts = n_experts
        self.n_experts_active = n_experts_active
        self.n_shared_experts = n_shared_experts
        self.moe_ffn_mult = moe_ffn_mult
        self.moe_aux_coef = moe_aux_coef
        self.moe_zloss_coef = moe_zloss_coef
        self.moe_norm_topk = moe_norm_topk
        self.emo_doc_pool = emo_doc_pool
        self.emo_bos_id = emo_bos_id