"""HevConfig: everything needed to rebuild the frozen backbone, the head, and the prompt protocol."""
from __future__ import annotations

try:
    from transformers import PreTrainedConfig
except ImportError:  # transformers < 5
    from transformers import PretrainedConfig as PreTrainedConfig

from .head import HevHeadConfig
from .schema import ALIASES


class HevConfig(PreTrainedConfig):
    model_type = "hev"

    def __init__(self, backbone_id: str = "google/diffusiongemma-26B-A4B-it", backbone_config: dict | None = None,
                 head_config: dict | None = None, canvas_length: int = 64, seed_mode: str = "pad",
                 slot_token_id: int = 0, encoder_norm: str = "", decoder_norm: str = "",
                 max_prompt_tokens: int = 1024, temperatures: list[float] | None = None,
                 n_heads: int = 1, member_temperatures: list[list[float]] | None = None,
                 aliases: list[str] | None = None, prompt_format_version: int = 1,
                 extra_layers: dict | None = None, **kwargs):
        self.backbone_id = backbone_id
        # Module names of the intermediate encoder/decoder layers whose unit-RMS outputs the features append to
        # the final-norm block (hev.features, --layers). Empty for a head trained on the final norms alone.
        extra_layers = extra_layers or {}
        self.extra_layers = {"encoder": list(extra_layers.get("encoder", [])),
                             "decoder": list(extra_layers.get("decoder", []))}
        self.backbone_config = backbone_config or {}
        self.head_config = head_config or HevHeadConfig().to_dict()
        self.canvas_length = canvas_length
        self.seed_mode = seed_mode
        self.slot_token_id = slot_token_id
        self.encoder_norm = encoder_norm
        self.decoder_norm = decoder_norm
        self.max_prompt_tokens = max_prompt_tokens
        self.temperatures = temperatures or [1.0, 1.0, 1.0]
        self.n_heads = int(n_heads)
        # One temperature triple per member; a single-head repo written before this field existed has none, so
        # its `temperatures` triple serves member 0.
        self.member_temperatures = member_temperatures or [list(self.temperatures)]
        self.aliases = aliases or list(ALIASES)
        # The hev.prompting format the head was trained with. Configs written before this field existed used 1.
        self.prompt_format_version = prompt_format_version
        kwargs.setdefault("auto_map", {"AutoConfig": "configuration_hev.HevConfig",
                                       "AutoModel": "modeling_hev.HevForDecisions"})
        super().__init__(**kwargs)