# coding=utf-8
# Configuration for Oleander (hybrid GDN2 / gated-GQA causal LMs).
# Auto-generated by the Oleander training notebook. Renaming this file means updating config.json -> auto_map
# ("AutoConfig": "configuration_oleander.OleanderConfig").
from __future__ import annotations

from transformers import PretrainedConfig


class OleanderConfig(PretrainedConfig):
    """Every architecture knob lives in the nested `oleander` dict written by the trainer (copy of CONFIG['model']).
    The flat keys (vocab_size, hidden_size, num_hidden_layers, ...) are duplicated so generic tooling can read them.
    """

    model_type = "oleander"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(self, vocab_size=32770, hidden_size=576, num_groups=7, group_pattern=None, mlp_hidden=1600,
                 max_seq_len=2048, norm_eps=1e-6, init_std=0.02, tie_embeddings=True, attn_res=True,
                 residual_dtype="bfloat16", gdn2=None, gqa=None, oleander=None, training=None, **kwargs):
        src = dict(oleander or {})
        self.vocab_size = int(src.get("vocab_size", vocab_size))
        self.hidden_size = int(src.get("hidden_size", hidden_size))
        self.num_groups = int(src.get("num_groups", num_groups))
        self.group_pattern = list(src.get("group_pattern", group_pattern or ["gdn2", "gdn2", "gdn2", "gqa"]))
        self.mlp_hidden = int(src.get("mlp_hidden", mlp_hidden))
        self.max_seq_len = int(src.get("max_seq_len", max_seq_len))
        self.norm_eps = float(src.get("norm_eps", norm_eps))
        self.init_std = float(src.get("init_std", init_std))
        self.tie_embeddings = bool(src.get("tie_embeddings", tie_embeddings))
        self.attn_res = bool(src.get("attn_res", attn_res))
        self.residual_dtype = str(src.get("residual_dtype", residual_dtype))
        self.gdn2 = dict(src.get("gdn2") or gdn2 or {})
        self.gqa = dict(src.get("gqa") or gqa or {})
        self.training = dict(training or kwargs.pop("training", {}) or {})
        # derived, so downstream code never has to recompute them
        self.num_hidden_layers = len(self.group_pattern) * self.num_groups
        # `oleander_layer_types` keeps our precise names; `layer_types` must stay inside the vocabulary that
        # transformers >= 5 validates against, so it gets the closest generic label.
        self.oleander_layer_types = [self.group_pattern[i % len(self.group_pattern)]
                                     for i in range(self.num_hidden_layers)]
        self.layer_types = ["linear_attention" if t == "gdn2" else "full_attention"
                            for t in self.oleander_layer_types]
        self.intermediate_size = self.mlp_hidden
        self.max_position_embeddings = self.max_seq_len
        self.rms_norm_eps = self.norm_eps
        self.num_attention_heads = int(self.gqa.get("q_heads", 0) or 0)
        self.num_key_value_heads = int(self.gqa.get("kv_heads", 0) or 0)
        self.head_dim = int(self.gqa.get("head_dim", 0) or 0)
        self.torch_dtype = kwargs.pop("torch_dtype", kwargs.pop("dtype", None))
        kwargs.setdefault("use_cache", False)
        _am = kwargs.pop("auto_map", None)        # keep it: save_pretrained() must preserve the remote-code hooks
        if _am:
            self.auto_map = _am
        kwargs.pop("tie_word_embeddings", None)   # we pass it explicitly
        kwargs.pop("layer_types", None)           # derived from group_pattern
        kwargs.pop("oleander_layer_types", None)  # derived from group_pattern
        kwargs.pop("model_type", None)
        super().__init__(tie_word_embeddings=self.tie_embeddings, **kwargs)

    def to_trainer_dict(self):
        """The dict the trainer expects (CONFIG['model']), for continued pre-training."""
        m = {"arch_name": "OleanderForCausalLM", "vocab_size": self.vocab_size, "hidden_size": self.hidden_size,
             "num_groups": self.num_groups, "group_pattern": list(self.group_pattern),
             "mlp_hidden": self.mlp_hidden, "max_seq_len": self.max_seq_len, "norm_eps": self.norm_eps,
             "init_std": self.init_std, "tie_embeddings": self.tie_embeddings, "attn_res": self.attn_res,
             "residual_dtype": self.residual_dtype, "gdn2": dict(self.gdn2), "gqa": dict(self.gqa)}
        return m