"""Configuration for Cherry Alpha."""

from __future__ import annotations

from transformers import PretrainedConfig


class CherryAlphaConfig(PretrainedConfig):
    model_type = "cherry_alpha"

    def __init__(
        self,
        vocab_size: int = 2048,
        hidden_size: int = 256,
        intermediate_size: int = 1488,
        num_hidden_layers: int = 3,
        num_attention_heads: int = 8,
        num_key_value_heads: int = 1,
        head_dim: int = 32,
        ngram_buckets: int = 156250,
        loop_schedule: tuple[int, ...] | list[int] = (0, 1, 1, 2),
        refresh_kernel_size: int = 9,
        max_position_embeddings: int = 4096,
        rope_theta: float = 100000.0,
        rms_norm_eps: float = 1e-5,
        initializer_range: float = 0.02,
        gradient_checkpointing: bool = True,
        tie_word_embeddings: bool = True,
        use_cache: bool = False,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        pad_token_id: int = 0,
        unk_token_id: int = 3,
        **kwargs,
    ):
        if hidden_size != num_attention_heads * head_dim:
            raise ValueError("hidden_size must equal num_attention_heads * head_dim")
        if num_attention_heads % num_key_value_heads:
            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
        if num_hidden_layers != 3:
            raise ValueError("Cherry Alpha requires exactly three physical layers")
        loop_schedule = tuple(int(index) for index in loop_schedule)
        if loop_schedule != (0, 1, 1, 2):
            raise ValueError("Cherry Alpha requires L1 -> L2 -> L2 -> L3")
        if ngram_buckets <= 0:
            raise ValueError("ngram_buckets must be positive")
        if refresh_kernel_size <= 0 or refresh_kernel_size % 2 == 0:
            raise ValueError("refresh_kernel_size must be a positive odd integer")

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.ngram_buckets = ngram_buckets
        self.loop_schedule = list(loop_schedule)
        self.refresh_kernel_size = refresh_kernel_size
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.rms_norm_eps = rms_norm_eps
        self.initializer_range = initializer_range
        self.gradient_checkpointing = gradient_checkpointing
        self.use_cache = False
        super().__init__(
            tie_word_embeddings=tie_word_embeddings,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            unk_token_id=unk_token_id,
            **kwargs,
        )

    def parameter_counts(self) -> dict[str, int]:
        hidden = self.hidden_size
        query_width = self.num_attention_heads * self.head_dim
        kv_width = self.num_key_value_heads * self.head_dim
        token_embedding = self.vocab_size * hidden
        ngram_embedding = self.ngram_buckets * hidden + 1
        attention = (
            hidden * query_width
            + 2 * hidden * kv_width
            + query_width * hidden
            + 2 * self.head_dim
        )
        refresh = (
            3 * hidden * hidden
            + hidden * self.refresh_kernel_size
            + 3 * hidden
            + 1
        )
        mlp = 3 * hidden * self.intermediate_size
        block_norms = 2 * hidden
        physical_layer = attention + refresh + mlp + block_norms
        neural_core = (
            token_embedding
            + self.num_hidden_layers * physical_layer
            + hidden
        )
        total = neural_core + ngram_embedding
        if not self.tie_word_embeddings:
            total += token_embedding
            neural_core += token_embedding
        return {
            "total": total,
            "neural_core": neural_core,
            "ngram": ngram_embedding,
            "token_embedding": token_embedding,
            "physical_layer": physical_layer,
            "physical_layers": self.num_hidden_layers,
            "effective_layer_passes": len(self.loop_schedule),
        }


CherryAlphaConfig.register_for_auto_class("AutoConfig")

__all__ = ["CherryAlphaConfig"]