"""Config for the memory-augmented Qwen3 (主路 Qwen3Attention ‖ 边路 GDN2)."""
from __future__ import annotations

from transformers.models.qwen3.configuration_qwen3 import Qwen3Config


class LiveMemConfig(Qwen3Config):
    """Qwen3 + per-attention-head GDN2 memory side-branch.

    All base Qwen3 fields are inherited unchanged. The `mem_*` fields configure
    the side branch and the two memory mechanisms (Design X / Design Y).
    """

    model_type = "livemem"

    def __init__(
        self,
        # which mechanism: "X" = 连续扫描 (continuous scan, write_mask=None);
        #                  "Y" = 门控读写解耦 (freeze gates on read tokens).
        memory_design: str = "Y",
        # layers that get a memory branch; None = all layers.
        mem_layers: list[int] | None = None,
        # GDN2 side-branch hyper-params. The legacy/default geometry is a
        # full-MHA copy of Qwen3 (8 KV heads repeated 4x -> 32 heads), per plan
        # 02 §1.3. Newer experiments can explicitly set mem_num_heads=8,
        # mem_num_v_heads=8, mem_expand_v=4 to keep similar state capacity with
        # fewer Q/K heads.
        mem_head_dim: int | None = None,
        mem_num_heads: int | None = None,
        mem_num_v_heads: int | None = None,
        mem_expand_v: float = 1.0,
        mem_conv_size: int = 4,
        mem_conv_bias: bool = False,
        mem_norm_eps: float | None = None,
        # zero-init the side o_proj so training starts ≈ original Qwen3 (plan T1).
        mem_o_proj_zero_init: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        if memory_design not in ("X", "Y"):
            raise ValueError(f"memory_design must be 'X' or 'Y', got {memory_design!r}")
        self.memory_design = memory_design
        self.mem_layers = mem_layers
        self.mem_head_dim = mem_head_dim if mem_head_dim is not None else self.head_dim
        self.mem_num_heads = (
            mem_num_heads if mem_num_heads is not None else self.num_attention_heads
        )
        self.mem_num_v_heads = (
            mem_num_v_heads if mem_num_v_heads is not None else self.mem_num_heads
        )
        self.mem_expand_v = mem_expand_v
        self.mem_conv_size = mem_conv_size
        self.mem_conv_bias = mem_conv_bias
        self.mem_norm_eps = mem_norm_eps if mem_norm_eps is not None else self.rms_norm_eps
        self.mem_o_proj_zero_init = mem_o_proj_zero_init

    @property
    def memory_layer_indices(self) -> list[int]:
        if self.mem_layers is None:
            return list(range(self.num_hidden_layers))
        return list(self.mem_layers)