"""Configuration for the complete MemT-1.0-Flash model (MIT)."""

import math
from copy import deepcopy

from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config


_DECLARED_SEMANTICS = {
    "memory_update_mode": "token",
    "memory_use_short_conv": False,
    "memory_head_mode": "mha",
    "memory_read_scale": "k_inv_sqrt",
    "tail_policy": "commit",
    "output_gate_type": "silu",
    "output_gate_softcap_enabled": True,
}


class MemTDeepseekV4Config(DeepseekV4Config):
    model_type = "memt_deepseek_v4"

    def __init__(self, memt=None, **kwargs):
        super().__init__(**kwargs)
        supplied = deepcopy(memt or {})
        defaults = {
            **_DECLARED_SEMANTICS,
            "layers": [i for i, kind in enumerate(self.layer_types)
                       if i >= 3 and kind == "compressed_sparse_attention"],
            "projection": "grouped_kv_low_rank",
            "num_memory_heads": self.num_attention_heads,
            "num_kv_heads": self.num_attention_heads,
            "head_dim": self.head_dim,
            "q_rank": self.q_lora_rank,
            "kv_groups": 8,
            "kv_rank": self.head_dim,
            "gate_rank": self.head_dim,
            "norm_eps": self.rms_norm_eps,
            "gate_latent_norm": False,
            "gdn_version": "stable_gdn2",
            "stable_gdn2_sqrt_erase_gate": True,
            "output_gate_bias": True,
            "output_gate_softcap_limit": 2.0,
            "chunk_tokens": 2048,
        }
        unknown = set(supplied) - set(defaults)
        if unknown:
            raise ValueError(f"Unknown MemT fields: {sorted(unknown)}")
        self.memt = defaults | supplied
        m = self.memt
        for name, expected in _DECLARED_SEMANTICS.items():
            if type(m[name]) is not type(expected) or m[name] != expected:
                raise ValueError(f"Unsupported memt.{name}: {m[name]!r}")
        if (m["projection"] != "grouped_kv_low_rank"
                or m["gdn_version"] != "stable_gdn2" or m["stable_gdn2_sqrt_erase_gate"] is not True
                or m["gate_latent_norm"]):
            raise ValueError("This release implements the trained grouped-KV stable sqrt-b model")
        if type(m["output_gate_bias"]) is not bool:
            raise ValueError("memt.output_gate_bias must be bool")
        limit = m["output_gate_softcap_limit"]
        if isinstance(limit, bool) or not isinstance(limit, (int, float)) or not math.isfinite(limit) or limit <= 0:
            raise ValueError("memt.output_gate_softcap_limit must be positive and finite")
        for name in ("num_memory_heads", "num_kv_heads", "head_dim", "q_rank", "kv_groups", "kv_rank", "gate_rank", "chunk_tokens"):
            if type(m[name]) is not int or m[name] <= 0:
                raise ValueError(f"memt.{name} must be a positive integer")
        if m["num_memory_heads"] % m["kv_groups"]:
            raise ValueError("MemT heads must be divisible by the number of KV groups")
        if m["num_kv_heads"] != m["num_memory_heads"]:
            raise ValueError("MemT KV heads must match memory heads after grouped projection")
        if m["num_memory_heads"] != self.num_attention_heads or m["head_dim"] != self.head_dim:
            raise ValueError("MemT output geometry must match the host attention")
        if len(set(m["layers"])) != len(m["layers"]) or any(
            not 0 <= i < self.num_hidden_layers or self.layer_types[i] != "compressed_sparse_attention"
            for i in m["layers"]
        ):
            raise ValueError("MemT layers must be distinct CSA layers")

    @classmethod
    def from_dict(cls, config_dict, **kwargs):
        result = super().from_dict(deepcopy(config_dict), **kwargs)
        config = result[0] if isinstance(result, tuple) else result
        # HF applies keyword overrides after __init__; validate those too.
        config.memt = cls(**config.to_dict()).memt
        return result


MemTDeepseekV4Config.register_for_auto_class()