"""Hugging Face configuration for ARLoop."""
from transformers import PretrainedConfig


class ARLoopConfig(PretrainedConfig):
    model_type = "arloop"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(self, vocab_size=49152, hidden_size=2048, intermediate_size=5632,
                 physical_layers=24, loop_count=3, num_attention_heads=16,
                 num_key_value_heads=4, max_position_embeddings=4096,
                 rope_theta=50000.0, rope_condense_ratio=1, rms_norm_eps=1e-5,
                 attnres_block_size=24, sw_blocks=3, qk_norm=True,
                 embedding_scale=1.0, logit_softcap=None, use_cache=True, **kwargs):
        kwargs.pop("num_hidden_layers", None)
        kwargs.pop("tie_word_embeddings", None)
        super().__init__(tie_word_embeddings=True, **kwargs)
        for key, value in dict(vocab_size=vocab_size, hidden_size=hidden_size,
                intermediate_size=intermediate_size, physical_layers=physical_layers,
                loop_count=loop_count, num_attention_heads=num_attention_heads,
                num_key_value_heads=num_key_value_heads,
                max_position_embeddings=max_position_embeddings, rope_theta=rope_theta,
                rope_condense_ratio=rope_condense_ratio, rms_norm_eps=rms_norm_eps,
                attnres_block_size=attnres_block_size, sw_blocks=sw_blocks,
                qk_norm=qk_norm, embedding_scale=embedding_scale,
                logit_softcap=logit_softcap, use_cache=use_cache).items():
            setattr(self, key, value)
        self.num_hidden_layers = physical_layers * loop_count
        self.head_dim = hidden_size // num_attention_heads
        if hidden_size % num_attention_heads or num_attention_heads % num_key_value_heads:
            raise ValueError("Invalid grouped-query attention dimensions")
        if self.head_dim % 2 or hidden_size < 32:
            raise ValueError("RoPE needs even head_dim; value gating needs hidden_size >= 32")
        if attnres_block_size < 2 or attnres_block_size % 2 or sw_blocks < 1:
            raise ValueError("Invalid residual block/window size")