"""HuggingFace configuration for Logos."""

from __future__ import annotations

from types import SimpleNamespace
from typing import Any, Dict

from transformers import PretrainedConfig


_NATIVE_DEFAULTS: Dict[str, Any] = {
    "vocab_size": 32000,
    "d_model": 512,
    "max_seq_len": 2048,
    "num_layers": 8,
    "num_heads": 8,
    "norm_eps": 1e-6,
    "d_ff": 1364,
    "use_moe": True,
    "num_shared_experts": 2,
    "num_sparse_experts": 64,
    "top_k": 6,
    "expert_d_ff": 256,
    "bias_update_rate": 0.01,
    "capacity_factor": 2.0,
    "router_logit_noise_std": 0.1,
    "router_logit_noise_decay_steps": 2000,
    "router_init_std": 0.002,
    "router_bias_error_clip": 1.0,
    "router_bias_clip": 1.0,
    "moe_aux_loss_weight": 1e-3,
    "moe_aux_loss_decay_steps": 2000,
    "lm_head_chunk_size": 0,
    "moe_diversity_factor": 0.0,
    "rope_base": 10000.0,
    "qk_norm": True,
    "partial_rope_dim": None,
    "attention_sink": True,
    "block_residual_isolate_softmax": False,
    "head_dim": 64,
    "conv_size": 4,
    "chunk_size": 64,
    "A_init_range": (1, 16),
    "expand": 2,
    "swa_window": 256,
    "swa_every": 4,
    "swa_offset": 3,
    "csa_compression": 4,
    "csa_top_k": 1024,
    "csa_indexer_heads": 4,
    "csa_indexer_dim": 32,
    "csa_indexer_loss_weight": 1.0,
    "hca_compression": 128,
    "compressed_query_dim": None,
    "compressed_head_dim": None,
    "compressed_rope": False,
    "indexer_rope": False,
    "attn_pattern": None,
    "num_entry_layers": 2,
    "num_body_layers": 4,
    "num_exit_layers": 2,
    "num_loops": 4,
    "entry_attn_pattern": None,
    "body_attn_pattern": None,
    "exit_attn_pattern": None,
    "entry_top_k": None,
    "exit_top_k": None,
    "gradient_checkpointing": False,
    "ckpt_granularity": "per-block",
}


class LogosConfig(PretrainedConfig):
    model_type = "logos"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(self, **kwargs: Any):
        native_values: Dict[str, Any] = {}
        for name, default in _NATIVE_DEFAULTS.items():
            native_values[name] = kwargs.pop(name, default)

        tokenizer_encoding = kwargs.pop("tokenizer_encoding", "cl100k_base")
        kwargs.setdefault("tie_word_embeddings", True)
        super().__init__(**kwargs)

        for name, value in native_values.items():
            setattr(self, name, value)
        self.tokenizer_encoding = tokenizer_encoding
        self.architectures = ["LogosForCausalLM"]
        self.auto_map = {
            "AutoConfig": "configuration_logos.LogosConfig",
            "AutoModelForCausalLM": "modeling_logos.LogosForCausalLM",
            "AutoTokenizer": "tokenization_logos.LogosTokenizer",
        }

    def to_native_config(self):
        data = {name: getattr(self, name) for name in _NATIVE_DEFAULTS}
        data["num_layers"] = (
            int(data["num_entry_layers"])
            + int(data["num_body_layers"])
            + int(data["num_exit_layers"])
        )
        return SimpleNamespace(**data)


__all__ = ["LogosConfig"]