"""Ethoric config. The short field names (n_layer, d_model, ...) are the real
ones; the aliases Transformers and lm-eval look for are set alongside them in
__init__.
"""
from __future__ import annotations

from transformers import PretrainedConfig

# `doc_masking`: Canon taps do not reach across an EOS. Set it false for an
# unmasked conv.
DEFAULT_CANON = {"enabled": True, "kernel": 4, "positions": ["A", "C"], "doc_masking": True}


class EthoricConfig(PretrainedConfig):
    """Decoder-only LM: RMSNorm, RoPE, grouped-query attention with QK-norm,
    SwiGLU, Canon convolutions, Peri-LN, tied embeddings.

    Two fields are easy to trip over. `attn_window_layers` lists the layers
    that attend to a short local window instead of the full context -- it is
    part of the trained architecture, and an empty list silently gives you a
    different model. `canon.doc_masking` only matters when a batch packs several
    documents behind EOS.
    """

    model_type = "ethoric"
    keys_to_ignore_at_inference = []

    def __init__(
        self,
        vocab_size: int = 5120,
        d_model: int = 128,
        n_layer: int = 12,
        n_head: int = 4,
        n_kv_head: int = 1,
        head_dim: int = 32,
        ffn_hidden: int = 352,
        max_seq_len: int = 512,
        rope_theta: float = 10000.0,
        norm_eps: float = 1e-5,
        qk_norm: str = "per_head",
        peri_ln: bool = True,
        attn_gate: str = "headwise",
        attn_window: int = 16,
        attn_window_layers: list | None = None,
        canon: dict | None = None,
        doc_masking: bool = True,
        tie_embeddings: bool = True,
        **kwargs,
    ):
        self.vocab_size = int(vocab_size)
        self.d_model = int(d_model)
        self.n_layer = int(n_layer)
        self.n_head = int(n_head)
        self.n_kv_head = int(n_kv_head)
        self.head_dim = int(head_dim)
        self.ffn_hidden = int(ffn_hidden)
        self.max_seq_len = int(max_seq_len)
        self.rope_theta = float(rope_theta)
        self.norm_eps = float(norm_eps)
        self.qk_norm = str(qk_norm)
        self.peri_ln = bool(peri_ln)
        self.attn_gate = str(attn_gate)
        self.attn_window = int(attn_window)
        self.attn_window_layers = list(attn_window_layers or [])
        self.canon = {**DEFAULT_CANON, **(canon or {})}
        self.doc_masking = bool(doc_masking)

        # Aliases the Transformers ecosystem looks for. Set as plain attributes
        # rather than through attribute_map so both names survive a
        # save_pretrained/from_pretrained round trip.
        self.hidden_size = self.d_model
        self.num_hidden_layers = self.n_layer
        self.num_attention_heads = self.n_head
        self.num_key_value_heads = self.n_kv_head
        self.intermediate_size = self.ffn_hidden
        self.max_position_embeddings = self.max_seq_len
        self.rms_norm_eps = self.norm_eps

        kwargs.setdefault("tie_word_embeddings", bool(tie_embeddings))
        kwargs.setdefault("bos_token_id", None)
        self.n_special = int(kwargs.pop("n_special", 2))
        kwargs.setdefault("eos_token_id", 0)
        kwargs.setdefault("pad_token_id", 1)
        super().__init__(**kwargs)

    @property
    def tie_embeddings(self) -> bool:
        return bool(self.tie_word_embeddings)

    @classmethod
    def from_pipeline(cls, cfg: dict) -> "EthoricConfig":
        """Build from a `model:` config block."""
        m = dict(cfg["model"])
        # Popped but unused: keeps a stray budget block out of the **m splat.
        m.pop("budget", None)
        # The boundary token comes from the tokenizer block: document masking
        # splits on it, and an adopted tokenizer need not put it at id 0.
        tok = cfg.get("tokenizer", {}) or {}
        eos = int(tok.get("eos_id", 0))
        n_special = len(tok.get("special_tokens", []) or []) or 2
        pad = next((i for i in range(n_special) if i != eos), None)
        return cls(
            tie_embeddings=m.pop("tie_embeddings", True),
            eos_token_id=eos,
            pad_token_id=pad,
            n_special=n_special,
            **m,
        )