from transformers import PretrainedConfig


class BananaMind21UnifiedConfig(PretrainedConfig):
    """BananaMind 2.1 Unified: three towers, a relay middle, one token out.

    A and C are the outer towers and each owns an output head.  B is the relay:
    no head, no solo loss, and the only path between A and C.
    """

    model_type = "bananamind21_unified"

    def __init__(
        self,
        vocab_size=8192,
        head_dim=64,
        max_position_embeddings=4096,
        rope_theta=100000.0,
        rms_norm_eps=1e-6,
        embed_width=384,
        hidden_size_a=256,
        num_hidden_layers_a=14,
        num_attention_heads_a=4,
        num_key_value_heads_a=1,
        intermediate_size_a=704,
        hidden_size_b=320,
        num_hidden_layers_b=5,
        num_attention_heads_b=5,
        num_key_value_heads_b=1,
        intermediate_size_b=960,
        hidden_size_c=384,
        num_hidden_layers_c=6,
        num_attention_heads_c=6,
        num_key_value_heads_c=2,
        intermediate_size_c=1024,
        a_read=(5, 9, 12),
        a_land=(7, 11, 14),
        c_read=(2, 4, 5),
        c_land=(3, 5, 6),
        b_land=(1, 3, 5),
        b_read=(2, 4, 5),
        gate_init=0.01,
        cut_bridges=False,
        relay_mode=None,
        use_single_tower=None,
        use_cache=True,
        tie_word_embeddings=False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.head_dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.rms_norm_eps = rms_norm_eps
        self.embed_width = embed_width

        self.hidden_size_a = hidden_size_a
        self.num_hidden_layers_a = num_hidden_layers_a
        self.num_attention_heads_a = num_attention_heads_a
        self.num_key_value_heads_a = num_key_value_heads_a
        self.intermediate_size_a = intermediate_size_a

        self.hidden_size_b = hidden_size_b
        self.num_hidden_layers_b = num_hidden_layers_b
        self.num_attention_heads_b = num_attention_heads_b
        self.num_key_value_heads_b = num_key_value_heads_b
        self.intermediate_size_b = intermediate_size_b

        self.hidden_size_c = hidden_size_c
        self.num_hidden_layers_c = num_hidden_layers_c
        self.num_attention_heads_c = num_attention_heads_c
        self.num_key_value_heads_c = num_key_value_heads_c
        self.intermediate_size_c = intermediate_size_c

        self.a_read = list(a_read)
        self.a_land = list(a_land)
        self.c_read = list(c_read)
        self.c_land = list(c_land)
        self.b_land = list(b_land)
        self.b_read = list(b_read)
        self.gate_init = gate_init

        # Ablation controls.  `relay_mode` is the general switch - see
        # `RELAY_MODES` in the modeling file for the full list:
        #
        #   full         three towers, all bridges, both heads mixed (default)
        #   cut_bridges  every bridge severed; A and C run standalone
        #   bypass_b     B's blocks skipped, bridges still live
        #   ab_only      tower C off; A and B run; head A alone
        #   cb_only      tower A off; C and B run; head C alone
        #   a_only       tower A alone, no bridges
        #   c_only       tower C alone, no bridges
        #
        # `use_single_tower="a"` / `"c"` is an alias for `a_only` / `c_only`, and
        # `cut_bridges=True` predates both and still selects `cut_bridges`.
        # Setting two of them to different modes is an error rather than a
        # silent precedence rule.
        self.cut_bridges = cut_bridges
        self.relay_mode = relay_mode
        self.use_single_tower = use_single_tower
        self.use_cache = use_cache

        # `hidden_size` is what generic HF tooling looks for
        self.hidden_size = hidden_size_c

        # Total attention layers over all three towers.  There is no single
        # "number of layers" in a relay model, but the KV cache needs one flat
        # index space, and generic HF tooling - `DynamicCache` above all - sizes
        # itself from `num_hidden_layers`.  The flat order is A, then B, then C:
        #
        #     A -> 0 .. n_a-1        B -> n_a .. n_a+n_b-1        C -> n_a+n_b ..
        #
        # `BananaMind21UnifiedForCausalLM` assigns the matching `layer_idx` to
        # every block, so the two must stay consistent.  Derived, never read
        # from a serialised config, so a stale value cannot desync the mapping.
        kwargs.pop("num_hidden_layers", None)
        self.num_hidden_layers = (
            num_hidden_layers_a + num_hidden_layers_b + num_hidden_layers_c
        )
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)