#!/usr/bin/env python3
"""BananaMind 2.1 Unified - three towers, relay middle, 35.0M params.

Three independent transformer stacks share one input embedding.  A and C are
the outer towers; each owns an output head and contributes to the mixture.  B is
the relay: it has no head, no solo loss term, and no gradient of its own except
what arrives through its four bridge directions.  There is no direct A<->C path,
so anything the two outer towers share about a token has to survive a trip
through B.

Exchange schedule
-----------------
Three rounds.  Each round reads the outer towers, lands the sum in B, lets B
process, then reads B and lands the result back in both outer towers.

    round | A read | C read | -> B  | B runs | B read | -> A  | -> C
      1   |   5    |   2    | pre-1 |  1-2   |   2    |   7   |   3
      2   |   9    |   4    | pre-3 |  3-4   |   4    |  11   |   5
      3   |  12    |   5    | pre-5 |   5    |   5    |  14   |   6

Bridge output is added to the residual *before* the receiving block rather than
after it, which is what guarantees B always has at least one full block between
taking a signal in and handing one back out.

Gates are per-channel and initialise to 0.01, not 0.  B is the only path between
A and C and it carries no solo loss, so a zero-init would leave the middle with
no gradient signal at all on step 0 and a real chance of never waking up.  0.01
is small enough that the towers still start as three near-independent models.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field

import torch
import torch.nn as nn
import torch.nn.functional as F


# --------------------------------------------------------------------------- #
# config
# --------------------------------------------------------------------------- #


@dataclass
class TowerSpec:
    hidden_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    intermediate_size: int


@dataclass
class RelayConfig:
    # inherited from Mini / 2.0 Unified, unchanged
    vocab_size: int = 8192
    head_dim: int = 64
    max_position_embeddings: int = 4096
    rope_theta: float = 100000.0
    rms_norm_eps: float = 1e-6

    # the shared embedding lives at tower C's width, so C reads it natively and
    # A and B each get one projection down.
    embed_width: int = 384

    tower_a: TowerSpec = field(
        default_factory=lambda: TowerSpec(
            hidden_size=256,
            num_hidden_layers=14,
            num_attention_heads=4,
            num_key_value_heads=1,
            intermediate_size=704,
        )
    )
    tower_b: TowerSpec = field(
        default_factory=lambda: TowerSpec(
            hidden_size=320,
            num_hidden_layers=5,
            num_attention_heads=5,
            num_key_value_heads=1,
            intermediate_size=960,
        )
    )
    tower_c: TowerSpec = field(
        default_factory=lambda: TowerSpec(
            hidden_size=384,
            num_hidden_layers=6,
            num_attention_heads=6,
            num_key_value_heads=2,
            intermediate_size=1024,
        )
    )

    # 1-indexed exchange schedule.  `*_read` is where a bridge takes a tower's
    # state; `*_land` is where the returning signal is added.  Biased deep on
    # the outer towers, because in the 2.0 run the deep bridges carried 4-25x
    # the gate magnitude of the shallow ones.
    a_read: tuple[int, ...] = (5, 9, 12)
    a_land: tuple[int, ...] = (7, 11, 14)
    c_read: tuple[int, ...] = (2, 4, 5)
    c_land: tuple[int, ...] = (3, 5, 6)
    b_land: tuple[int, ...] = (1, 3, 5)
    b_read: tuple[int, ...] = (2, 4, 5)

    gate_init: float = 0.01

    # untied on purpose: A and C each own an output space, so either one still
    # decodes by itself once the bridges come off.
    tie_word_embeddings: bool = False

    @property
    def n_rounds(self) -> int:
        return len(self.a_read)

    def validate(self) -> None:
        if self.embed_width != self.tower_c.hidden_size:
            raise ValueError("shared embedding width must match tower C width")
        n = self.n_rounds
        for name in ("a_read", "a_land", "c_read", "c_land", "b_land", "b_read"):
            if len(getattr(self, name)) != n:
                raise ValueError(f"{name} must have {n} entries")

        # Every tower is visited strictly left to right, and every bridge reads
        # a state that already exists.  If this passes, the forward pass below
        # is a valid topological order.
        for read, land, spec, tag in (
            (self.a_read, self.a_land, self.tower_a, "A"),
            (self.c_read, self.c_land, self.tower_c, "C"),
        ):
            cursor = 0
            for k in range(n):
                # `==` is legal: the previous round's signal is added before that
                # layer runs, so reading right after it is still causal.
                if not cursor <= read[k]:
                    raise ValueError(f"tower {tag} round {k}: read layer {read[k]} is behind cursor {cursor}")
                if not read[k] < land[k]:
                    raise ValueError(f"tower {tag} round {k}: land {land[k]} must follow read {read[k]}")
                cursor = land[k]
            if max(land) > spec.num_hidden_layers or min(read) < 1:
                raise ValueError(f"tower {tag} schedule leaves the stack")

        cursor = 0
        for k in range(n):
            if not cursor < self.b_land[k]:
                raise ValueError(f"tower B round {k}: land {self.b_land[k]} is not ahead of {cursor}")
            if not self.b_land[k] <= self.b_read[k]:
                raise ValueError(f"tower B round {k}: read {self.b_read[k]} precedes land {self.b_land[k]}")
            cursor = self.b_read[k]
        if max(self.b_read) > self.tower_b.num_hidden_layers or min(self.b_land) < 1:
            raise ValueError("tower B schedule leaves the stack")


# --------------------------------------------------------------------------- #
# primitives (identical to train_mini.py / 2.0 Unified)
# --------------------------------------------------------------------------- #


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_float = x.float()
        rms = torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x_float * rms * self.weight.float()).type_as(x)


def build_rope_inv_freq(head_dim: int, theta: float = 100000.0) -> torch.Tensor:
    return 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))


def precompute_freqs_cis(head_dim: int, seq_len: int, theta: float = 100000.0) -> torch.Tensor:
    freqs = torch.outer(torch.arange(seq_len, dtype=torch.float32), build_rope_inv_freq(head_dim, theta))
    return torch.polar(torch.ones_like(freqs), freqs)


def apply_rotary_emb(
    q: torch.Tensor, k: torch.Tensor, freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    q_complex = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
    k_complex = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(0)
    q_out = torch.view_as_real(q_complex * freqs_cis).flatten(-2)
    k_out = torch.view_as_real(k_complex * freqs_cis).flatten(-2)
    return q_out.type_as(q), k_out.type_as(k)


class TowerAttention(nn.Module):
    def __init__(self, spec: TowerSpec, head_dim: int, rms_norm_eps: float):
        super().__init__()
        self.n_head = spec.num_attention_heads
        self.n_kv_heads = spec.num_key_value_heads
        self.head_dim = head_dim
        self.n_rep = self.n_head // self.n_kv_heads

        self.q_proj = nn.Linear(spec.hidden_size, self.n_head * head_dim, bias=False)
        self.k_proj = nn.Linear(spec.hidden_size, self.n_kv_heads * head_dim, bias=False)
        self.v_proj = nn.Linear(spec.hidden_size, self.n_kv_heads * head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_head * head_dim, spec.hidden_size, bias=False)
        self.o_proj.NANOGPT_SCALE_INIT = 1

        self.q_norm = RMSNorm(head_dim, eps=rms_norm_eps)
        self.k_norm = RMSNorm(head_dim, eps=rms_norm_eps)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        bsz, seq_len, _ = x.size()

        q = self.q_proj(x).view(bsz, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)

        q = self.q_norm(q)
        k = self.k_norm(k)
        q, k = apply_rotary_emb(q, k, freqs_cis)

        # .contiguous() is required for torch.compile: without it the expanded
        # view reaches the flash-attention backward with strides Inductor
        # rejects.  Pure layout change, no effect on the math.
        k = (
            k.unsqueeze(2)
            .expand(bsz, self.n_kv_heads, self.n_rep, seq_len, self.head_dim)
            .reshape(bsz, self.n_head, seq_len, self.head_dim)
            .contiguous()
        )
        v = (
            v.unsqueeze(2)
            .expand(bsz, self.n_kv_heads, self.n_rep, seq_len, self.head_dim)
            .reshape(bsz, self.n_head, seq_len, self.head_dim)
            .contiguous()
        )

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(bsz, seq_len, self.n_head * self.head_dim)
        return self.o_proj(y)


class TowerSwiGLUMLP(nn.Module):
    def __init__(self, spec: TowerSpec):
        super().__init__()
        self.w_gate = nn.Linear(spec.hidden_size, spec.intermediate_size, bias=False)
        self.w_up = nn.Linear(spec.hidden_size, spec.intermediate_size, bias=False)
        self.w_down = nn.Linear(spec.intermediate_size, spec.hidden_size, bias=False)
        self.w_down.NANOGPT_SCALE_INIT = 1

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class TowerBlock(nn.Module):
    """A plain Mini block.  Nothing in here knows about bridges."""

    def __init__(self, spec: TowerSpec, head_dim: int, rms_norm_eps: float):
        super().__init__()
        self.ln_1 = RMSNorm(spec.hidden_size, eps=rms_norm_eps)
        self.attn = TowerAttention(spec, head_dim, rms_norm_eps)
        self.ln_2 = RMSNorm(spec.hidden_size, eps=rms_norm_eps)
        self.mlp = TowerSwiGLUMLP(spec)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x), freqs_cis)
        x = x + self.mlp(self.ln_2(x))
        return x


# --------------------------------------------------------------------------- #
# bridge
# --------------------------------------------------------------------------- #


class Edge(nn.Module):
    """One direction of one exchange point.

    W is a plain linear, g a per-channel gate.  The gate starts small but not
    zero so the relay gets traffic on step 0.
    """

    def __init__(self, dim_in: int, dim_out: int, gate_init: float):
        super().__init__()
        self.w = nn.Linear(dim_in, dim_out, bias=False)
        self.g = nn.Parameter(torch.full((dim_out,), float(gate_init)))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.g * self.w(x)


# --------------------------------------------------------------------------- #
# model
# --------------------------------------------------------------------------- #


class BananaMindRelay(nn.Module):
    def __init__(self, config: RelayConfig):
        super().__init__()
        config.validate()
        self.config = config
        a, b, c = config.tower_a, config.tower_b, config.tower_c

        self.wte = nn.Embedding(config.vocab_size, config.embed_width)
        self.in_proj_a = nn.Linear(config.embed_width, a.hidden_size, bias=False)
        self.in_proj_b = nn.Linear(config.embed_width, b.hidden_size, bias=False)
        # tower C reads the shared embedding natively; no projection

        self.blocks_a = nn.ModuleList(
            [TowerBlock(a, config.head_dim, config.rms_norm_eps) for _ in range(a.num_hidden_layers)]
        )
        self.blocks_b = nn.ModuleList(
            [TowerBlock(b, config.head_dim, config.rms_norm_eps) for _ in range(b.num_hidden_layers)]
        )
        self.blocks_c = nn.ModuleList(
            [TowerBlock(c, config.head_dim, config.rms_norm_eps) for _ in range(c.num_hidden_layers)]
        )

        n = config.n_rounds
        gi = config.gate_init
        self.edges_a2b = nn.ModuleList([Edge(a.hidden_size, b.hidden_size, gi) for _ in range(n)])
        self.edges_c2b = nn.ModuleList([Edge(c.hidden_size, b.hidden_size, gi) for _ in range(n)])
        self.edges_b2a = nn.ModuleList([Edge(b.hidden_size, a.hidden_size, gi) for _ in range(n)])
        self.edges_b2c = nn.ModuleList([Edge(b.hidden_size, c.hidden_size, gi) for _ in range(n)])

        # B has no final norm because it has no head: its only outputs are the
        # bridge reads, which are normalised by the receiving tower's blocks.
        self.ln_f_a = RMSNorm(a.hidden_size, eps=config.rms_norm_eps)
        self.ln_f_c = RMSNorm(c.hidden_size, eps=config.rms_norm_eps)

        self.lm_head_a = nn.Linear(a.hidden_size, config.vocab_size, bias=False)
        self.lm_head_c = nn.Linear(c.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head_c.weight = self.wte.weight

        # still a 2-way mixer: it reads the two towers that have heads
        self.mix_head = nn.Linear(a.hidden_size + c.hidden_size, 1, bias=True)

        self._embd_scale = math.sqrt(config.embed_width)
        self._freqs_cis_cache: torch.Tensor | None = None

        self._init_all()

    # ---- init -------------------------------------------------------------- #

    def _init_module(self, module: nn.Module, depth: int) -> None:
        std = 0.02
        if hasattr(module, "NANOGPT_SCALE_INIT"):
            std *= 2 * depth**-0.5
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def _init_all(self) -> None:
        cfg = self.config
        depth_a = cfg.tower_a.num_hidden_layers
        depth_b = cfg.tower_b.num_hidden_layers
        depth_c = cfg.tower_c.num_hidden_layers

        for blocks, depth in (
            (self.blocks_a, depth_a),
            (self.blocks_b, depth_b),
            (self.blocks_c, depth_c),
        ):
            for m in blocks.modules():
                self._init_module(m, depth)

        for m in (
            self.wte,
            self.in_proj_a,
            self.in_proj_b,
            self.lm_head_a,
            self.lm_head_c,
            self.mix_head,
        ):
            self._init_module(m, depth_a)

        for edges, depth in (
            (self.edges_a2b, depth_b),
            (self.edges_c2b, depth_b),
            (self.edges_b2a, depth_a),
            (self.edges_b2c, depth_c),
        ):
            for edge in edges:
                self._init_module(edge.w, depth)
                torch.nn.init.constant_(edge.g, cfg.gate_init)

        # start the mixer at alpha = 0.5 so neither outer tower is favoured
        torch.nn.init.zeros_(self.mix_head.weight)
        torch.nn.init.zeros_(self.mix_head.bias)

    # ---- rope cache -------------------------------------------------------- #

    def _get_freqs_cis(self, seq_len: int, device: torch.device) -> torch.Tensor:
        cache = self._freqs_cis_cache
        if cache is None or cache.device != device or cache.size(0) < seq_len:
            cache = precompute_freqs_cis(
                self.config.head_dim, seq_len, self.config.rope_theta
            ).to(device)
            self._freqs_cis_cache = cache
        return cache[:seq_len]

    # ---- forward ----------------------------------------------------------- #

    @staticmethod
    def _advance(
        blocks: nn.ModuleList,
        x: torch.Tensor,
        cursor: int,
        target: int,
        freqs_cis: torch.Tensor,
        pending: dict[int, torch.Tensor],
    ) -> tuple[torch.Tensor, int]:
        """Run `blocks` from `cursor` up to and including 1-indexed `target`.

        A pending bridge value keyed by a 1-indexed layer is added to the
        residual immediately before that layer runs.
        """
        while cursor < target:
            layer = cursor + 1
            arriving = pending.pop(layer, None)
            if arriving is not None:
                x = x + arriving
            x = blocks[cursor](x, freqs_cis)
            cursor += 1
        return x, cursor

    def forward(
        self,
        input_ids: torch.Tensor,
        cut_bridges: bool = False,
        collect_diagnostics: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict]:
        """Returns (h_a, h_c, mix_logit, diagnostics).

        With `cut_bridges=True` the three stacks run completely independently:
        A and C become ordinary transformers and B becomes dead weight that
        influences nothing.
        """
        cfg = self.config
        _, seq_len = input_ids.size()
        freqs_cis = self._get_freqs_cis(seq_len, input_ids.device)

        embedded = self.wte(input_ids) * self._embd_scale
        x_a = self.in_proj_a(embedded)
        x_b = self.in_proj_b(embedded)
        x_c = embedded

        pend_a: dict[int, torch.Tensor] = {}
        pend_b: dict[int, torch.Tensor] = {}
        pend_c: dict[int, torch.Tensor] = {}
        cur_a = cur_b = cur_c = 0
        diagnostics: dict = {}

        for k in range(cfg.n_rounds):
            # 1. outer towers walk to their read points, applying whatever the
            #    previous round handed them on the way.
            x_a, cur_a = self._advance(self.blocks_a, x_a, cur_a, cfg.a_read[k], freqs_cis, pend_a)
            x_c, cur_c = self._advance(self.blocks_c, x_c, cur_c, cfg.c_read[k], freqs_cis, pend_c)

            # 2. both outer signals land in the relay at the same layer.
            if not cut_bridges:
                pend_b[cfg.b_land[k]] = self.edges_a2b[k](x_a) + self.edges_c2b[k](x_c)

            # 3. the relay processes, then hands one signal back to each side.
            x_b, cur_b = self._advance(self.blocks_b, x_b, cur_b, cfg.b_read[k], freqs_cis, pend_b)

            if not cut_bridges:
                pend_a[cfg.a_land[k]] = self.edges_b2a[k](x_b)
                pend_c[cfg.c_land[k]] = self.edges_b2c[k](x_b)

            if collect_diagnostics:
                with torch.no_grad():
                    fa, fb, fc = x_a.float(), x_b.float(), x_c.float()
                    diagnostics[f"cos_a2b/round{k}"] = float(
                        F.cosine_similarity(self.edges_a2b[k].w(fa), fb, dim=-1).mean()
                    )
                    diagnostics[f"cos_c2b/round{k}"] = float(
                        F.cosine_similarity(self.edges_c2b[k].w(fc), fb, dim=-1).mean()
                    )
                    diagnostics[f"cos_b2a/round{k}"] = float(
                        F.cosine_similarity(self.edges_b2a[k].w(fb), fa, dim=-1).mean()
                    )
                    diagnostics[f"cos_b2c/round{k}"] = float(
                        F.cosine_similarity(self.edges_b2c[k].w(fb), fc, dim=-1).mean()
                    )
                    # how much of A survives the trip to C and back through B
                    diagnostics[f"rms_a/round{k}"] = float(fa.pow(2).mean().sqrt())
                    diagnostics[f"rms_b/round{k}"] = float(fb.pow(2).mean().sqrt())
                    diagnostics[f"rms_c/round{k}"] = float(fc.pow(2).mean().sqrt())

        # 4. finish every stack
        x_a, cur_a = self._advance(
            self.blocks_a, x_a, cur_a, cfg.tower_a.num_hidden_layers, freqs_cis, pend_a
        )
        x_c, cur_c = self._advance(
            self.blocks_c, x_c, cur_c, cfg.tower_c.num_hidden_layers, freqs_cis, pend_c
        )
        # B's tail runs only so its parameters see gradient in later rounds; with
        # the default schedule B is already finished here.
        x_b, cur_b = self._advance(
            self.blocks_b, x_b, cur_b, cfg.tower_b.num_hidden_layers, freqs_cis, pend_b
        )

        if pend_a or pend_b or pend_c:
            raise RuntimeError(f"bridge values were never consumed: {sorted(pend_a) + sorted(pend_b) + sorted(pend_c)}")

        h_a = self.ln_f_a(x_a)
        h_c = self.ln_f_c(x_c)
        mix_logit = self.mix_head(torch.cat([h_a, h_c], dim=-1)).squeeze(-1)
        return h_a, h_c, mix_logit, diagnostics

    # ---- diagnostics ------------------------------------------------------- #

    @torch.no_grad()
    def gate_stats(self) -> dict[str, float]:
        """Per-round, per-direction gate magnitudes."""
        cfg = self.config
        stats: dict[str, float] = {}
        groups = (
            ("a2b", self.edges_a2b, cfg.b_land),
            ("c2b", self.edges_c2b, cfg.b_land),
            ("b2a", self.edges_b2a, cfg.a_land),
            ("b2c", self.edges_b2c, cfg.c_land),
        )
        for name, edges, layers in groups:
            for i, edge in enumerate(edges):
                g = edge.g.detach().float()
                stats[f"gate_{name}/L{layers[i]}_absmean"] = float(g.abs().mean())
                stats[f"gate_{name}/L{layers[i]}_absmax"] = float(g.abs().max())
                stats[f"gate_{name}/L{layers[i]}_rms"] = float(g.pow(2).mean().sqrt())
        return stats

    # ---- parameter accounting ---------------------------------------------- #

    def param_breakdown(self) -> dict[str, int]:
        def n(module: nn.Module) -> int:
            return sum(p.numel() for p in module.parameters())

        bridges = sum(n(e) for e in (self.edges_a2b, self.edges_c2b, self.edges_b2a, self.edges_b2c))
        return {
            "shared_embedding": n(self.wte),
            "in_proj_384_to_256": n(self.in_proj_a),
            "in_proj_384_to_320": n(self.in_proj_b),
            "tower_a_layers": n(self.blocks_a),
            "tower_b_layers": n(self.blocks_b),
            "tower_c_layers": n(self.blocks_c),
            "bridges": bridges,
            "ln_f_a": n(self.ln_f_a),
            "ln_f_c": n(self.ln_f_c),
            "lm_head_a": n(self.lm_head_a),
            "lm_head_c": n(self.lm_head_c),
            "mix_head": n(self.mix_head),
            "total": sum(p.numel() for p in self.parameters()),
        }

    def matmul_parameters(self) -> int:
        """Total excluding norms, gates and the mixer.

        This is the number the architecture is specified against: everything
        that is a weight matrix, and nothing that is a per-channel scalar.
        """
        skip = set()
        for module in self.modules():
            if isinstance(module, RMSNorm):
                skip.add(id(module.weight))
        for edges in (self.edges_a2b, self.edges_c2b, self.edges_b2a, self.edges_b2c):
            for edge in edges:
                skip.add(id(edge.g))
        for p in self.mix_head.parameters():
            skip.add(id(p))
        return sum(p.numel() for p in self.parameters() if id(p) not in skip)


# --------------------------------------------------------------------------- #
# loss
# --------------------------------------------------------------------------- #


def _loss_chunk(
    h_a: torch.Tensor,
    h_c: torch.Tensor,
    mix_logit: torch.Tensor,
    targets: torch.Tensor,
    w_a: torch.Tensor,
    w_c: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """One chunk of positions.  Returns summed NLLs and summed squared logits."""
    logits_a = F.linear(h_a, w_a).float()
    logits_c = F.linear(h_c, w_c).float()

    log_p_a = F.log_softmax(logits_a, dim=-1)
    log_p_c = F.log_softmax(logits_c, dim=-1)

    # Mixing happens in probability space, computed in log space:
    #   log(a*p_A + (1-a)*p_C) = logaddexp(log a + log p_A, log(1-a) + log p_C)
    log_alpha = F.logsigmoid(mix_logit).unsqueeze(-1)
    log_one_minus_alpha = F.logsigmoid(-mix_logit).unsqueeze(-1)
    log_p_mix = torch.logaddexp(log_alpha + log_p_a, log_one_minus_alpha + log_p_c)

    idx = targets.unsqueeze(-1)
    nll_mix = -log_p_mix.gather(-1, idx).squeeze(-1).sum()
    nll_a = -log_p_a.gather(-1, idx).squeeze(-1).sum()
    nll_c = -log_p_c.gather(-1, idx).squeeze(-1).sum()
    sq_logits = logits_a.pow(2).sum() + logits_c.pow(2).sum()
    return nll_mix, nll_a, nll_c, sq_logits


def relay_loss(
    model: BananaMindRelay,
    h_a: torch.Tensor,
    h_c: torch.Tensor,
    mix_logit: torch.Tensor,
    targets: torch.Tensor,
    solo_lambda: float,
    z_loss_coeff: float = 0.0,
    chunk_size: int = 16384,
) -> dict[str, torch.Tensor]:
    """L = L_mix + lambda * (L_A + L_C).

    There is no L_B term.  B has no head to compute one with, and giving it one
    would turn the relay into a third predictor instead of a channel.  B trains
    entirely on gradient arriving through its four bridge directions.

    Positions are processed in chunks under gradient checkpointing: three
    vocab-sized float tensors per position do not fit at 4096 context otherwise.
    """
    w_a = model.lm_head_a.weight
    w_c = model.lm_head_c.weight
    vocab = w_a.size(0)

    flat_a = h_a.reshape(-1, h_a.size(-1))
    flat_c = h_c.reshape(-1, h_c.size(-1))
    flat_z = mix_logit.reshape(-1)
    flat_y = targets.reshape(-1)
    n_pos = flat_y.numel()

    if chunk_size <= 0 or chunk_size >= n_pos:
        bounds = [(0, n_pos)]
    else:
        bounds = [(i, min(i + chunk_size, n_pos)) for i in range(0, n_pos, chunk_size)]

    mix_sum = flat_a.new_zeros((), dtype=torch.float32)
    a_sum = flat_a.new_zeros((), dtype=torch.float32)
    c_sum = flat_a.new_zeros((), dtype=torch.float32)
    sq_sum = flat_a.new_zeros((), dtype=torch.float32)

    use_ckpt = torch.is_grad_enabled() and len(bounds) > 1
    for lo, hi in bounds:
        args = (flat_a[lo:hi], flat_c[lo:hi], flat_z[lo:hi], flat_y[lo:hi], w_a, w_c)
        if use_ckpt:
            part = torch.utils.checkpoint.checkpoint(_loss_chunk, *args, use_reentrant=False)
        else:
            part = _loss_chunk(*args)
        mix_sum = mix_sum + part[0]
        a_sum = a_sum + part[1]
        c_sum = c_sum + part[2]
        sq_sum = sq_sum + part[3]

    l_mix = mix_sum / n_pos
    l_a = a_sum / n_pos
    l_c = c_sum / n_pos
    z_loss = sq_sum / (2.0 * n_pos * vocab)

    loss = l_mix + solo_lambda * (l_a + l_c) + z_loss_coeff * z_loss
    return {"loss": loss, "l_mix": l_mix, "l_a": l_a, "l_c": l_c, "z_loss": z_loss}


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


def estimate_training_flops_per_token(config: RelayConfig, param_count: int, seq_len: int) -> int:
    dense_flops = 6 * param_count
    attention_flops = 0
    for spec in (config.tower_a, config.tower_b, config.tower_c):
        attention_flops += (
            12 * spec.num_hidden_layers * spec.num_attention_heads * config.head_dim * seq_len
        )
    return dense_flops + attention_flops