"""Fugue: masked-diffusion language models on a bidirectional Gated DeltaNet-2 hybrid.

Fugue keeps Senryu-2's block (4-branch Gated Residual stream, SiTU-GLU,
full-rank sigmoid output gates) and makes every token mixer bidirectional:

- Bidirectional GDN-2 reads each segment left to right and right to left over
  one set of q/k/v, erase/write and output projections. Each direction has its
  own short convolutions and low-rank decay (ViG). The backward scan runs on
  the segment reversed in place, so both scans share ``cu_seqlens``, and it
  reads its state before the current token (Hydra), so each token's own key
  and value enter its output once.
- Bidirectional attention is grouped-query softmax attention with per-head
  QK-norm, partial RoPE and a sigmoid output gate.

Masked diffusion (MDLM, LLaDA): each window draws a mask ratio t and every
token becomes ``<|mask|>`` with probability t. The model predicts the
originals; ``<|mask|>`` itself always gets probability zero. There is no time
input; the mask count carries it. The trainer picks the per-token weight: 1/t
for the ELBO, or 1 (unweighted masked cross-entropy). Evaluation always uses
the ELBO (docs/architecture.md, "Fugue").

Decoding runs equal-length canvases as a batch rather than packed, with no
host sync, so a CUDA graph can capture each call: ``canvas_logits`` scores
picked positions. It can also cache everything outside the block being
denoised (Fast-dLLM's DualCache): ``refresh_block`` runs the whole canvas once
and keeps, per layer, each scan's state and conv window where it enters the
block and the attention keys and values; ``block_logits`` then runs only the
block's tokens against them. Text outside the block keeps its features from
the refresh.

Fugue-1 (``fugue_block_size`` > 0) writes fixed-size blocks left to right
instead, over a causal clean stream whose cache is exact; ``tiny_gdn.fugue_block``
holds that mode's training forward and decoding, and the whole-sequence
methods here refuse it.
"""

from __future__ import annotations

import itertools
import math
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from fla.modules import FusedRMSNormGated, ShortConvolution
from fla.modules.rotary import rotary_embedding
from fla.ops.gdn2 import chunk_gdn2
from fla.ops.gdn2.fused_recurrent import fused_recurrent_gdn2_fwd
from safetensors.torch import save_model
from torch import nn
from torch.utils.checkpoint import checkpoint

from tiny_gdn.config import FUGUE_MODEL_TYPE, TinyGDNConfig
from tiny_gdn.fused_ops import weighted_linear_cross_entropy
from tiny_gdn.haiku_layers import SiTUGLU
from tiny_gdn.model import TinyGDNOutput
from tiny_gdn.nn_common import RMSNorm
from tiny_gdn.senryu2 import (
    GatedResidual,
    RotaryTables,
    mark_adamw_only,
    mark_gain,
    mark_muon_row_blocks,
)

IGNORE_INDEX = -100
# Smallest mask ratio drawn in training; the 1/t weight stays below 1,000.
MIN_MASK_RATIO = 1e-3


# --------------------------------------------------------------------------- layout


@dataclass(frozen=True)
class SegmentLayout:
    """Segments packed back to back in one row, and index maps within each segment."""

    cu_seqlens: torch.Tensor
    # Each position's mirror inside its segment; applying it twice is the identity.
    reverse_index: torch.Tensor
    # The next and previous position in the same segment; segment ends map to themselves.
    next_index: torch.Tensor
    previous_index: torch.Tensor
    segment_start: torch.Tensor


@dataclass(frozen=True)
class ScanEdge:
    """One scan direction of a GDN-2 layer where it enters the block: its state and q/k/v conv windows."""

    state: torch.Tensor
    convs: tuple[torch.Tensor, torch.Tensor, torch.Tensor]


@dataclass(frozen=True)
class GatedDeltaNetEdges:
    forward: ScanEdge
    backward: ScanEdge


@dataclass(frozen=True)
class AttentionContext:
    """Rotated keys and values of the whole canvas; each block step overwrites the block's rows."""

    keys: torch.Tensor
    values: torch.Tensor


LayerCache = GatedDeltaNetEdges | AttentionContext


@dataclass
class BlockCapture:
    """What a decoding forward keeps for later block-only steps.

    Per scan direction, in scan order, ``lead`` marks the tokens the scan reads
    before it reaches the block (the prefix forwards, the suffix backwards) and
    ``lead_count`` counts them. Both are tensors, so one CUDA graph serves every block.
    """

    lead: tuple[torch.Tensor, torch.Tensor]
    lead_count: tuple[torch.Tensor, torch.Tensor]
    edges: dict[tuple[int, int], ScanEdge] = field(default_factory=dict)
    attention: dict[int, AttentionContext] = field(default_factory=dict)

    @classmethod
    def around(cls, positions: torch.Tensor, length: int) -> BlockCapture:
        """The capture for the block at consecutive ``positions`` of ``length``-token rows."""

        scan = torch.arange(length, device=positions.device)
        before = positions[0]
        after = length - 1 - positions[-1]
        return cls(lead=(scan < before, scan < after), lead_count=(before, after))


@dataclass(frozen=True)
class BlockCache:
    """Everything outside the block ``[start, stop)`` as the last full forward saw it.

    ``offsets`` (the block start per row) and ``positions`` are tensors, so a
    CUDA graph captured against one cache serves every block of the canvas
    once ``copy_`` has loaded that block's cache into it.
    """

    layers: tuple[LayerCache, ...]
    rotary: tuple[torch.Tensor, torch.Tensor]
    offsets: torch.Tensor
    positions: torch.Tensor

    @property
    def rows(self) -> int:
        return self.offsets.shape[0]

    @property
    def width(self) -> int:
        return self.positions.shape[0]

    def tensors(self) -> Iterator[torch.Tensor]:
        for layer in self.layers:
            if isinstance(layer, AttentionContext):
                yield from (layer.keys, layer.values)
            else:
                for edge in (layer.forward, layer.backward):
                    yield edge.state
                    yield from edge.convs
        yield from self.rotary
        yield from (self.offsets, self.positions)

    def clone(self) -> BlockCache:
        def edge(source: ScanEdge) -> ScanEdge:
            q, k, v = (conv.clone() for conv in source.convs)
            return ScanEdge(state=source.state.clone(), convs=(q, k, v))

        return BlockCache(
            layers=tuple(
                AttentionContext(keys=layer.keys.clone(), values=layer.values.clone())
                if isinstance(layer, AttentionContext)
                else GatedDeltaNetEdges(forward=edge(layer.forward), backward=edge(layer.backward))
                for layer in self.layers
            ),
            rotary=(self.rotary[0].clone(), self.rotary[1].clone()),
            offsets=self.offsets.clone(),
            positions=self.positions.clone(),
        )

    def copy_(self, source: BlockCache) -> None:
        for target, value in zip(self.tensors(), source.tensors(), strict=True):
            target.copy_(value)


@dataclass
class FugueContext:
    """Per-forward layout shared by every block.

    Training and evaluation pack all rows into one ``[1, length]`` sequence
    split by ``layout.cu_seqlens``. Decoding runs ``[rows, length]`` canvases
    that are each one whole segment as a batch (``packed=False``), so no kernel
    needs the boundaries on the host.
    """

    length: int
    layout: SegmentLayout
    # Set when every segment has this length, so attention runs as one batched call.
    segment_length: int | None
    rotary: tuple[torch.Tensor, torch.Tensor]
    capture: BlockCapture | None = None
    rows: int = 1
    packed: bool = True

    @property
    def cu_seqlens(self) -> torch.Tensor | None:
        return self.layout.cu_seqlens if self.packed else None


def whole_row_layout(length: int, device: torch.device) -> SegmentLayout:
    """The layout of rows that are each one segment of ``length``, built without a host sync."""

    positions = torch.arange(length, device=device)
    return SegmentLayout(
        cu_seqlens=torch.arange(2, device=device, dtype=torch.int32) * length,
        reverse_index=positions.flip(0),
        next_index=(positions + 1).clamp(max=length - 1),
        previous_index=(positions - 1).clamp(min=0),
        segment_start=positions.eq(0),
    )


def segment_layout(lengths: torch.Tensor) -> SegmentLayout:
    lengths = lengths.to(torch.long)
    cu_seqlens = F.pad(lengths.cumsum(0), (1, 0))
    total = int(cu_seqlens[-1])
    positions = torch.arange(total, device=lengths.device)
    segment = torch.repeat_interleave(torch.arange(lengths.numel(), device=lengths.device), lengths)
    offset = positions - cu_seqlens[segment]
    segment_start = offset.eq(0)
    segment_end = offset.eq(lengths[segment] - 1)
    return SegmentLayout(
        cu_seqlens=cu_seqlens.to(torch.int32),
        reverse_index=cu_seqlens[segment] + cu_seqlens[segment + 1] - 1 - positions,
        next_index=torch.where(segment_end, positions, positions + 1),
        previous_index=torch.where(segment_start, positions, positions - 1),
        segment_start=segment_start,
    )


# ------------------------------------------------------------------- token mixers


class BidirectionalGatedDeltaNet2(nn.Module):
    """GDN-2 in both directions over shared projections, averaged before the output gate.

    S_t = (I - k_t (b_t ⊙ k_t)ᵀ) Diag(exp g_t) S_{t-1} + k_t (w_t ⊙ v_t)ᵀ runs
    left to right and right to left; o_t = ½ q_tᵀ (S→_t + S←_{t+1}). Index 0 of
    each per-direction module or parameter is the forward scan, index 1 the
    backward one.
    """

    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
        super().__init__()
        hidden = config.hidden_size
        self.layer_index = layer_index
        self.num_heads = config.linear_num_heads
        self.num_v_heads = config.linear_num_value_heads
        self.head_k_dim = config.linear_head_dim
        self.head_v_dim = int(config.linear_head_dim * config.linear_expand_v)
        self.key_dim = self.num_heads * self.head_k_dim
        self.value_dim = self.num_v_heads * self.head_v_dim
        self.allow_neg_eigval = config.allow_negative_eigenvalues
        kernel = config.linear_conv_kernel_dim

        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
        self.q_conv1d = nn.ModuleList(
            ShortConvolution(self.key_dim, kernel, bias=False, activation="silu") for _ in range(2)
        )
        self.k_conv1d = nn.ModuleList(
            ShortConvolution(self.key_dim, kernel, bias=False, activation="silu") for _ in range(2)
        )
        self.v_conv1d = nn.ModuleList(
            ShortConvolution(self.value_dim, kernel, bias=False, activation="silu") for _ in range(2)
        )
        self.f_proj = nn.ModuleList(
            nn.Sequential(
                nn.Linear(hidden, self.head_v_dim, bias=False),
                nn.Linear(self.head_v_dim, self.key_dim, bias=False),
            )
            for _ in range(2)
        )
        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
        self.A_log = nn.Parameter(torch.empty(2, self.num_heads, dtype=torch.float32).uniform_(1, 16).log())
        dt = torch.exp(
            torch.rand(2, self.key_dim, dtype=torch.float32) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
        ).clamp(min=1e-4)
        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
        if config.linear_full_rank_gate:
            self.g_proj: nn.Module = nn.Linear(hidden, self.value_dim, bias=True)
            mark_adamw_only(self.g_proj.weight)
        else:
            self.g_proj = nn.Sequential(
                nn.Linear(hidden, self.head_v_dim, bias=False),
                nn.Linear(self.head_v_dim, self.value_dim, bias=True),
            )
            mark_adamw_only(*(layer.weight for layer in self.g_proj))
        self.o_norm = FusedRMSNormGated(self.head_v_dim, activation="sigmoid", eps=config.rms_norm_eps)
        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)

        mark_gain(self.A_log, self.dt_bias)
        mark_adamw_only(*(layer.weight for decay in self.f_proj for layer in decay))
        for projection in (self.q_proj, self.k_proj, self.b_proj):
            mark_muon_row_blocks(projection.weight, self.num_heads)
        for projection in (self.v_proj, self.w_proj):
            mark_muon_row_blocks(projection.weight, self.num_v_heads)

    def convolutions(self, direction: int) -> tuple[ShortConvolution, ShortConvolution, ShortConvolution]:
        return self.q_conv1d[direction], self.k_conv1d[direction], self.v_conv1d[direction]

    def heads(
        self,
        direction: int,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        f: torch.Tensor,
        b: torch.Tensor,
        w: torch.Tensor,
    ) -> tuple[torch.Tensor, ...]:
        """Per-head ``(q, k, v, g, b, w)`` for the scan kernels, decay applied and key heads repeated."""

        g = F.softplus(f + self.dt_bias[direction])
        q, k, g, b = (rearrange(tensor, "... (h d) -> ... h d", d=self.head_k_dim) for tensor in (q, k, g, b))
        v, w = (rearrange(tensor, "... (h d) -> ... h d", d=self.head_v_dim) for tensor in (v, w))
        g = -self.A_log[direction].float().exp().unsqueeze(-1) * g
        if self.num_v_heads > self.num_heads:
            groups = self.num_v_heads // self.num_heads
            q, k, g, b = (repeat(tensor, "... h d -> ... (h g) d", g=groups) for tensor in (q, k, g, b))
        return q, k, v, g, b, w

    def scan(
        self,
        direction: int,
        hidden_states: torch.Tensor,
        projected: tuple[torch.Tensor, ...],
        layout: SegmentLayout,
        cu_seqlens: torch.Tensor | None,
        capture: BlockCapture | None = None,
    ) -> torch.Tensor:
        """One direction over every segment of ``layout``; ``cu_seqlens`` None scans rows that are one segment each."""

        f = self.f_proj[direction](hidden_states).float()
        q, k, v, f, b, w = (*projected[:3], f, *projected[3:])
        if direction == 1:
            q, k, v, f, b, w = (tensor.index_select(1, layout.reverse_index) for tensor in (q, k, v, f, b, w))
        unconvolved = (q, k, v)
        q_conv, k_conv, v_conv = self.convolutions(direction)
        q, _ = q_conv(x=q, cu_seqlens=cu_seqlens)
        k, _ = k_conv(x=k, cu_seqlens=cu_seqlens)
        v, _ = v_conv(x=v, cu_seqlens=cu_seqlens)
        if direction == 1:
            # Each query reads the state one step early; its output moves back a step below.
            q = q.index_select(1, layout.next_index)
        q, k, v, g, b, w = self.heads(direction, q, k, v, f, b, w)
        if capture is not None:
            self._capture_edge(direction, unconvolved, (k, v, g, b, w), capture)
        out, _ = chunk_gdn2(
            q=q,
            k=k,
            v=v,
            g=g,
            b=b,
            w=w,
            initial_state=None,
            output_final_state=False,
            use_qk_l2norm_in_kernel=True,
            cu_seqlens=cu_seqlens,
        )
        if direction == 0:
            return out
        out = out.index_select(1, layout.previous_index).masked_fill(layout.segment_start[None, :, None, None], 0.0)
        return out.index_select(1, layout.reverse_index)

    def _capture_edge(
        self,
        direction: int,
        unconvolved: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
        scanned: tuple[torch.Tensor, ...],
        capture: BlockCapture,
    ) -> None:
        """Keep this direction's state and conv windows after the tokens it reads before the block.

        The scan reruns over the whole row with zero key and decay past the
        lead, which leaves the state unchanged, so its shape is the same for
        every block.
        """

        lead, count = capture.lead[direction], capture.lead_count[direction]
        width = self.q_conv1d[direction].kernel_size[0]
        # A conv window holds the last `width` inputs, oldest first, zeros before the row starts.
        slots = count - width + torch.arange(width, device=lead.device)
        present = slots.ge(0)[None, :, None]
        q_window, k_window, v_window = (
            (tensor.index_select(1, slots.clamp(min=0)) * present).transpose(1, 2).contiguous()
            for tensor in unconvolved
        )
        k, v, g, b, w = scanned
        keep = lead[None, :, None, None]
        k = k * keep
        _, state = chunk_gdn2(
            q=k, k=k, v=v, g=g * keep, b=b, w=w, output_final_state=True, use_qk_l2norm_in_kernel=True
        )
        capture.edges[(self.layer_index, direction)] = ScanEdge(state=state, convs=(q_window, k_window, v_window))

    def scan_block(
        self,
        direction: int,
        hidden_states: torch.Tensor,
        projected: tuple[torch.Tensor, ...],
        edge: ScanEdge,
    ) -> torch.Tensor:
        """``scan`` over ``[rows, width]`` block tokens, entering from ``edge``."""

        f = self.f_proj[direction](hidden_states).float()
        q, k, v, f, b, w = (*projected[:3], f, *projected[3:])
        if direction == 1:
            q, k, v, f, b, w = (tensor.flip(1) for tensor in (q, k, v, f, b, w))
        # The kernels update conv windows in place; the cache must stay as the refresh left it.
        q_conv, k_conv, v_conv = self.convolutions(direction)
        q, _ = q_conv(x=q, cache=edge.convs[0].clone())
        k, _ = k_conv(x=k, cache=edge.convs[1].clone())
        v, _ = v_conv(x=v, cache=edge.convs[2].clone())
        q, k, v, g, b, w = self.heads(direction, q, k, v, f, b, w)
        if direction == 1:
            # A leading slot with zero key, value and decay passes the edge state
            # through, so query i reads the state before token i (Hydra).
            q = F.pad(q, (0, 0, 0, 0, 0, 1))
            k, v, g, b, w = (F.pad(tensor, (0, 0, 0, 0, 1, 0)) for tensor in (k, v, g, b, w))
        out, _ = fused_recurrent_gdn2_fwd(
            q=q.contiguous(),
            k=k.contiguous(),
            v=v.contiguous(),
            g=g.contiguous(),
            b=b.contiguous(),
            w=w.contiguous(),
            initial_state=edge.state,
            use_qk_l2norm_in_kernel=True,
        )
        if direction == 0:
            return out
        return out[:, :-1].flip(1)

    def project(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, ...]:
        b = self.b_proj(hidden_states).sigmoid()
        if self.allow_neg_eigval:
            b = b * 2.0
        return (
            self.q_proj(hidden_states),
            self.k_proj(hidden_states),
            self.v_proj(hidden_states),
            b,
            self.w_proj(hidden_states).sigmoid(),
        )

    def gated_output(self, out: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
        gate = rearrange(self.g_proj(hidden_states), "... (h d) -> ... h d", d=self.head_v_dim)
        out = self.o_norm(out, gate)
        return self.o_proj(rearrange(out, "b t h d -> b t (h d)"))

    def forward(self, hidden_states: torch.Tensor, context: FugueContext) -> torch.Tensor:
        projected = self.project(hidden_states)
        out = 0.5 * (
            self.scan(0, hidden_states, projected, context.layout, context.cu_seqlens, context.capture)
            + self.scan(1, hidden_states, projected, context.layout, context.cu_seqlens, context.capture)
        )
        return self.gated_output(out, hidden_states)

    def forward_block(self, hidden_states: torch.Tensor, cache: LayerCache, block: BlockCache) -> torch.Tensor:
        """``forward`` for ``[rows, width]`` block tokens between the cached edges."""

        if not isinstance(cache, GatedDeltaNetEdges):
            raise TypeError(f"layer {self.layer_index} needs GDN-2 edges, got {type(cache).__name__}")
        projected = self.project(hidden_states)
        out = 0.5 * (
            self.scan_block(0, hidden_states, projected, cache.forward)
            + self.scan_block(1, hidden_states, projected, cache.backward)
        )
        return self.gated_output(out, hidden_states)


class BidirectionalAttention(nn.Module):
    """Grouped-query softmax attention over each whole segment.

    Per-head zero-centered RMSNorm on q and k, partial RoPE, and a full-rank
    sigmoid output gate, as in Senryu-2's Parallax layers without the causal mask.
    """

    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
        super().__init__()
        self.layer_index = layer_index
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.attention_head_dim
        self.scale = self.head_dim**-0.5
        query_width = self.num_heads * self.head_dim
        kv_width = self.num_kv_heads * self.head_dim
        hidden = config.hidden_size
        self.q_proj = nn.Linear(hidden, query_width, bias=False)
        self.k_proj = nn.Linear(hidden, kv_width, bias=False)
        self.v_proj = nn.Linear(hidden, kv_width, bias=False)
        self.g_proj = nn.Linear(hidden, query_width, bias=True)
        self.o_proj = nn.Linear(query_width, hidden, bias=False)
        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        mark_muon_row_blocks(self.q_proj.weight, self.num_heads)
        for projection in (self.k_proj, self.v_proj):
            mark_muon_row_blocks(projection.weight, self.num_kv_heads)
        mark_adamw_only(self.g_proj.weight)

    def attend(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        """``[n, T, H, D]`` queries over ``[n, T, H_kv, D]`` keys/values, no mask."""

        out = F.scaled_dot_product_attention(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
            scale=self.scale,
            enable_gqa=self.num_kv_heads != self.num_heads,
        )
        return out.transpose(1, 2)

    def forward(self, hidden_states: torch.Tensor, context: FugueContext) -> torch.Tensor:
        rows, length = hidden_states.shape[:2]
        q = self.q_norm(self.q_proj(hidden_states).view(rows, length, self.num_heads, self.head_dim))
        k = self.k_norm(self.k_proj(hidden_states).view(rows, length, self.num_kv_heads, self.head_dim))
        v = self.v_proj(hidden_states).view(rows, length, self.num_kv_heads, self.head_dim)
        cosine, sine = context.rotary
        q, k = (
            rotary_embedding(tensor, cosine, sine, seqlen_offsets=0, cu_seqlens=context.cu_seqlens)
            for tensor in (q, k)
        )
        if context.capture is not None:
            context.capture.attention[self.layer_index] = AttentionContext(keys=k, values=v)
        if context.segment_length is not None:
            segments = rows * length // context.segment_length
            out = self.attend(
                q.view(segments, context.segment_length, self.num_heads, self.head_dim),
                k.view(segments, context.segment_length, self.num_kv_heads, self.head_dim),
                v.view(segments, context.segment_length, self.num_kv_heads, self.head_dim),
            ).reshape(rows, length, self.num_heads, self.head_dim)
        else:
            bounds = context.layout.cu_seqlens.tolist()
            out = torch.cat(
                [
                    self.attend(q[:, start:stop], k[:, start:stop], v[:, start:stop])
                    for start, stop in itertools.pairwise(bounds)
                ],
                dim=1,
            )
        gate = torch.sigmoid(self.g_proj(hidden_states)).view(rows, length, self.num_heads, self.head_dim)
        return self.o_proj((out * gate).reshape(rows, length, self.num_heads * self.head_dim))

    def forward_block(self, hidden_states: torch.Tensor, cache: LayerCache, block: BlockCache) -> torch.Tensor:
        """``[rows, width]`` block queries over the cached canvas, with the block's own keys and values."""

        if not isinstance(cache, AttentionContext):
            raise TypeError(f"layer {self.layer_index} needs attention context, got {type(cache).__name__}")
        rows, width = hidden_states.shape[:2]
        q = self.q_norm(self.q_proj(hidden_states).view(rows, width, self.num_heads, self.head_dim))
        k = self.k_norm(self.k_proj(hidden_states).view(rows, width, self.num_kv_heads, self.head_dim))
        v = self.v_proj(hidden_states).view(rows, width, self.num_kv_heads, self.head_dim)
        cosine, sine = block.rotary
        q, k = (rotary_embedding(tensor, cosine, sine, seqlen_offsets=block.offsets) for tensor in (q, k))
        cache.keys.index_copy_(1, block.positions, k)
        cache.values.index_copy_(1, block.positions, v)
        out = self.attend(q, cache.keys, cache.values)
        gate = torch.sigmoid(self.g_proj(hidden_states)).view(rows, width, self.num_heads, self.head_dim)
        return self.o_proj((out * gate).reshape(rows, width, self.num_heads * self.head_dim))


# ------------------------------------------------------------------------- blocks


class FugueBlock(nn.Module):
    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
        super().__init__()
        self.layer_index = layer_index
        self.layer_type = config.layer_types[layer_index]
        residual = (config.residual_branches, config.hidden_size, config.residual_gate_rank, config.rms_norm_eps)
        self.mixer_residual = GatedResidual(*residual, write=True)
        self.mlp_residual = GatedResidual(*residual, write=True)
        if self.layer_type == "bigdn2":
            self.token_mixer: nn.Module = BidirectionalGatedDeltaNet2(config, layer_index)
        elif self.layer_type == "bidirectional":
            self.token_mixer = BidirectionalAttention(config, layer_index)
        else:
            raise ValueError(f"Fugue has no {self.layer_type} layers")
        self.mlp = SiTUGLU(config)
        mark_muon_row_blocks(self.mlp.gate_up_proj.weight, 2)

    def mix(self, residual: torch.Tensor, mixer: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor:
        hidden_states, scales = self.mixer_residual.read(residual)
        residual = GatedResidual.write(residual, scales, mixer(hidden_states).reshape(hidden_states.shape))
        hidden_states, scales = self.mlp_residual.read(residual)
        return GatedResidual.write(residual, scales, self.mlp(hidden_states))

    def forward(self, residual: torch.Tensor, context: FugueContext) -> torch.Tensor:
        return self.mix(
            residual, lambda hidden: self.token_mixer(hidden.view(context.rows, context.length, -1), context)
        )

    def forward_block(self, residual: torch.Tensor, cache: LayerCache, block: BlockCache) -> torch.Tensor:
        return self.mix(
            residual,
            lambda hidden: self.token_mixer.forward_block(hidden.view(block.rows, block.width, -1), cache, block),
        )


# -------------------------------------------------------------------------- noise


def stratified_mask_ratios(
    count: int,
    *,
    generator: torch.Generator,
    device: torch.device | str,
    minimum: float = MIN_MASK_RATIO,
) -> torch.Tensor:
    """``count`` low-discrepancy mask ratios in ``[minimum, 1]``: one uniform offset, evenly spaced."""

    offset = torch.rand((), generator=generator, device=device)
    strata = (offset + torch.arange(count, device=device, dtype=torch.float32) / count) % 1.0
    order = torch.randperm(count, generator=generator, device=device)
    return minimum + (1.0 - minimum) * strata[order]


LOSS_WEIGHTINGS = ("elbo", "unweighted")


def mask_tokens(
    tokens: torch.Tensor,
    valid: torch.Tensor,
    ratios: torch.Tensor,
    *,
    mask_token_id: int,
    generator: torch.Generator,
    weighting: str = "elbo",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward noise: ``(noisy ids, labels with -100 where kept, per-token weights)`` for ``[B, T]`` tokens.

    Weights are 1/t for the ELBO and 1 for unweighted masked cross-entropy.
    """

    if weighting not in LOSS_WEIGHTINGS:
        raise ValueError(f"weighting must be one of: {', '.join(LOSS_WEIGHTINGS)}")
    draws = torch.rand(tokens.shape, generator=generator, device=tokens.device)
    masked = (draws < ratios[:, None]) & valid
    noisy = torch.where(masked, torch.full_like(tokens, mask_token_id), tokens)
    labels = torch.where(masked, tokens, torch.full_like(tokens, IGNORE_INDEX))
    row_weights = 1.0 / ratios if weighting == "elbo" else torch.ones_like(ratios)
    return noisy, labels, row_weights[:, None].expand_as(tokens)


def mask_exact_counts(
    tokens: torch.Tensor,
    valid: torch.Tensor,
    counts: torch.Tensor,
    *,
    mask_token_id: int,
    generator: torch.Generator,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Mask exactly ``counts[b]`` valid positions of each row, uniformly: ``(noisy ids, labels)``.

    With ``l ~ U{1..L}`` masked and weight L/l, this is LLaDA's lower-variance
    ELBO estimator; per token it is the mean CE over the masked positions.
    """

    draws = torch.rand(tokens.shape, generator=generator, device=tokens.device).masked_fill(~valid, 2.0)
    ranks = draws.argsort(dim=1).argsort(dim=1)
    masked = ranks < counts[:, None]
    noisy = torch.where(masked, torch.full_like(tokens, mask_token_id), tokens)
    labels = torch.where(masked, tokens, torch.full_like(tokens, IGNORE_INDEX))
    return noisy, labels


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


class FugueForMaskedDiffusion(nn.Module):
    """Bidirectional GDN-2 hybrid that predicts the tokens hidden behind ``<|mask|>``."""

    def __init__(self, config: TinyGDNConfig) -> None:
        super().__init__()
        if config.model_type != FUGUE_MODEL_TYPE or config.mask_token_id is None:
            raise ValueError("FugueForMaskedDiffusion needs a Fugue config with a mask_token_id")
        self.config = config
        self.mask_token_id = config.mask_token_id
        # Fugue-1's block size; 0 for whole-sequence denoising (Rondo).
        self.block_size = config.fugue_block_size
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(FugueBlock(config, index) for index in range(config.num_hidden_layers))
        self.final_residual = GatedResidual(
            config.residual_branches,
            config.hidden_size,
            config.residual_gate_rank,
            config.rms_norm_eps,
            write=False,
        )
        self.rotary = RotaryTables(config.rotary_dim, config.rope_theta)
        self.checkpointed_layers: frozenset[int] = frozenset()
        self.apply(self._initialize_module)
        self._initialize_residual_projections()

    # ---------------------------------------------------------------- setup
    def _initialize_module(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)

    def _initialize_residual_projections(self) -> None:
        residual_std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers)
        for layer in self.layers:
            nn.init.normal_(layer.token_mixer.o_proj.weight, mean=0.0, std=residual_std)
            nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=residual_std)

    def enable_gradient_checkpointing(self, enabled: bool = True, *, layers: int | None = None) -> None:
        """Recompute block activations in backward; ``layers`` picks GDN-2 blocks first."""

        if not enabled:
            if layers is not None:
                raise ValueError("layers needs gradient checkpointing enabled")
            self.checkpointed_layers = frozenset()
            return
        count = len(self.layers) if layers is None else layers
        if not 1 <= count <= len(self.layers):
            raise ValueError(f"layers must be between 1 and {len(self.layers)}")
        cheapest_replay_first = sorted(
            range(len(self.layers)),
            key=lambda index: self.layers[index].layer_type != "bigdn2",
        )
        self.checkpointed_layers = frozenset(cheapest_replay_first[:count])

    def describe_gradient_checkpointing(self) -> str:
        if not self.checkpointed_layers:
            return "off"
        layer_types = sorted({self.layers[index].layer_type for index in self.checkpointed_layers})
        return f"{len(self.checkpointed_layers)}/{len(self.layers)} blocks ({', '.join(layer_types)})"

    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Logits over the vocabulary; ``<|mask|>`` is never a prediction (logit -inf)."""

        logits = F.linear(hidden_states, self.embed_tokens.weight)
        logits[..., self.mask_token_id] = float("-inf")
        return logits

    # -------------------------------------------------------------- forward
    def _require_whole_sequence(self) -> None:
        if self.block_size:
            raise ValueError(
                "this Fugue model denoises blocks over a causal clean stream "
                f"(fugue_block_size={self.block_size}); train and decode it with tiny_gdn.fugue_block"
            )

    def _run_block(self, index: int, residual: torch.Tensor, context: FugueContext) -> torch.Tensor:
        layer = self.layers[index]
        if self.training and index in self.checkpointed_layers:
            return checkpoint(layer, residual, context, use_reentrant=False)
        return layer(residual, context)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        *,
        labels: torch.Tensor | None = None,
        label_weights: torch.Tensor | None = None,
        loss_normalizer: float | None = None,
        z_loss_coefficient: float = 0.0,
        return_logits: bool = True,
        logits_mask: torch.Tensor | None = None,
        return_hidden_states: bool = False,
    ) -> TinyGDNOutput:
        """Score masked positions of ``[B, T]`` rows; each row is one bidirectional segment.

        ``labels`` holds the original token at scored positions and -100
        elsewhere; ``loss = Σ w·CE / loss_normalizer`` (default: the number of
        scored positions). ``logits_mask`` limits logits to its true positions
        and returns them flattened ``[selected, vocab]`` in row-major order.
        """

        self._require_whole_sequence()
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        batch, length = input_ids.shape
        if length > self.config.max_position_embeddings:
            raise ValueError("Input exceeds max_position_embeddings")
        device = input_ids.device
        if attention_mask is None or bool(attention_mask.all()):
            index = None
            lengths = torch.full((batch,), length, dtype=torch.long, device=device)
            segment_length: int | None = length
        else:
            if attention_mask.shape != input_ids.shape:
                raise ValueError("attention_mask must match input_ids")
            valid = attention_mask.reshape(-1).bool()
            index = valid.nonzero().flatten()
            lengths = attention_mask.bool().sum(dim=1)
            if bool((lengths == 0).any()):
                raise ValueError("every row needs at least one unmasked token")
            segment_length = None

        def pack(tensor: torch.Tensor) -> torch.Tensor:
            flat = tensor.reshape(batch * length, *tensor.shape[2:])
            return flat if index is None else flat.index_select(0, index)

        tokens = pack(input_ids)
        embeddings = self.embed_tokens(tokens)
        context = FugueContext(
            length=tokens.numel(),
            layout=segment_layout(lengths),
            segment_length=segment_length,
            rotary=self.rotary.get(tokens.numel(), device, embeddings.dtype),
        )
        residual = embeddings.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
        for layer_index in range(len(self.layers)):
            residual = self._run_block(layer_index, residual, context)
        hidden_states, _ = self.final_residual.read(residual)

        total_loss = main_loss = z_loss = None
        if labels is not None:
            if labels.shape != input_ids.shape:
                raise ValueError("labels must match input_ids")
            targets = pack(labels)
            scored = targets.ne(IGNORE_INDEX).nonzero().flatten()
            weights = (
                torch.ones(scored.numel(), dtype=torch.float32, device=device)
                if label_weights is None
                else pack(label_weights).index_select(0, scored).float()
            )
            normalizer = float(scored.numel()) if loss_normalizer is None else float(loss_normalizer)
            if scored.numel() == 0:
                total_loss = hidden_states.sum() * 0.0
                main_loss = total_loss.detach()
                z_loss = total_loss.detach()
            else:
                objective, losses, log_partitions = weighted_linear_cross_entropy(
                    hidden_states.index_select(0, scored),
                    self.embed_tokens.weight,
                    targets.index_select(0, scored),
                    weights,
                    z_loss_coefficient=z_loss_coefficient,
                    excluded_token_id=self.mask_token_id,
                )
                total_loss = objective / max(normalizer, 1.0)
                main_loss = (weights * losses).sum() / max(normalizer, 1.0)
                z_loss = log_partitions.square().mean()

        logits = None
        if return_logits:
            if logits_mask is not None:
                selected = pack(logits_mask).nonzero().flatten()
                logits = self.project_to_vocabulary(hidden_states.index_select(0, selected))
            elif index is None:
                logits = self.project_to_vocabulary(hidden_states).view(batch, length, -1)
            else:
                full = hidden_states.new_zeros(batch * length, hidden_states.shape[-1])
                full.index_copy_(0, index, hidden_states)
                logits = self.project_to_vocabulary(full).view(batch, length, -1)
        return TinyGDNOutput(
            loss=total_loss,
            logits=logits,
            main_loss=main_loss,
            mtp_loss=None,
            z_loss=z_loss if z_loss_coefficient > 0.0 else None,
            hidden_states=hidden_states if return_hidden_states else None,
        )

    # ------------------------------------------------------------- decoding
    def _decode(self, canvas: torch.Tensor, capture: BlockCapture | None) -> torch.Tensor:
        """Final hidden states ``[rows, T, hidden]`` of unpadded canvases, run as a batch without host syncs."""

        self._require_whole_sequence()
        rows, length = canvas.shape
        if length > self.config.max_position_embeddings:
            raise ValueError("Input exceeds max_position_embeddings")
        embeddings = self.embed_tokens(canvas.reshape(-1))
        context = FugueContext(
            length=length,
            layout=whole_row_layout(length, canvas.device),
            segment_length=length,
            rotary=self.rotary.get(length, canvas.device, embeddings.dtype),
            capture=capture,
            rows=rows,
            packed=False,
        )
        residual = embeddings.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
        for layer in self.layers:
            residual = layer(residual, context)
        hidden_states, _ = self.final_residual.read(residual)
        return hidden_states.view(rows, length, -1)

    def canvas_logits(self, canvas: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        """Logits ``[rows, picks, vocab]`` at ``positions`` ``[rows, picks]`` of ``[rows, T]`` canvases."""

        hidden_states = self._decode(canvas, None)
        picked = hidden_states.gather(1, positions.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1]))
        return self.project_to_vocabulary(picked)

    def refresh_block(self, canvas: torch.Tensor, positions: torch.Tensor) -> tuple[torch.Tensor, BlockCache]:
        """Full forward of ``[rows, T]`` canvases: logits ``[rows, width, vocab]`` for the block and its cache.

        ``positions`` holds the block's consecutive canvas positions, after at
        least one token. The cache keeps ``positions`` itself, so a graph that
        refreshes from a static tensor caches a block that later calls can move.
        """

        rows, length = canvas.shape
        capture = BlockCapture.around(positions, length)
        hidden_states = self._decode(canvas, capture)
        logits = self.project_to_vocabulary(hidden_states.index_select(1, positions))
        layers: list[LayerCache] = []
        for index, layer in enumerate(self.layers):
            if layer.layer_type == "bigdn2":
                layers.append(GatedDeltaNetEdges(forward=capture.edges[(index, 0)], backward=capture.edges[(index, 1)]))
            else:
                layers.append(capture.attention[index])
        cache = BlockCache(
            layers=tuple(layers),
            rotary=self.rotary.get(length, canvas.device, self.embed_tokens.weight.dtype),
            offsets=positions[:1].to(torch.int32).repeat(rows),
            positions=positions,
        )
        return logits, cache

    def block_logits(self, block_ids: torch.Tensor, cache: BlockCache) -> torch.Tensor:
        """Logits ``[rows, width, vocab]`` for the block's current ``[rows, width]`` tokens against ``cache``."""

        self._require_whole_sequence()
        if block_ids.shape != (cache.rows, cache.width):
            raise ValueError("block_ids must match the cached block")
        embeddings = self.embed_tokens(block_ids.reshape(-1))
        residual = embeddings.unsqueeze(1).expand(-1, self.config.residual_branches, -1)
        for layer, layer_cache in zip(self.layers, cache.layers, strict=True):
            residual = layer.forward_block(residual, layer_cache, cache)
        hidden_states, _ = self.final_residual.read(residual)
        return self.project_to_vocabulary(hidden_states).view(cache.rows, cache.width, -1)

    # ------------------------------------------------------------ reporting
    def parameter_report(self) -> dict[str, int]:
        total = sum(parameter.numel() for parameter in self.parameters())
        embeddings = self.embed_tokens.weight.numel()
        return {
            "deployable_core": total,
            "training_total": total,
            "embedding": embeddings,
            "mtp_auxiliary": 0,
            "non_embedding_core": total - embeddings,
            "ngram_table": 0,
        }

    def save_checkpoint(self, output_dir: Path) -> None:
        output_dir.mkdir(parents=True, exist_ok=True)
        self.config.save_json(output_dir / "config.json")
        save_model(self, output_dir / "model.safetensors")

    def extra_repr(self) -> str:
        return f"core_parameters={self.parameter_report()['deployable_core']:,}"

    def get_architecture_metadata(self) -> dict[str, Any]:
        if self.block_size:
            objective = [
                (
                    f"block diffusion over {self.block_size}-token blocks: a causal clean stream with a next-token "
                    "loss, and a noisy stream denoising each block against the clean prefix (exact block cache)"
                ),
                "token editing: corrupted visible tokens are trained back to the originals",
                "3:1 Gated DeltaNet-2 + attention hybrid, bidirectional inside the active block",
                (
                    "GDN-2 noisy reads: the clean forward state where the block starts, plus forward and "
                    "backward scans inside the block (exclusive backward read, Hydra)"
                ),
            ]
        else:
            objective = [
                "masked diffusion (MDLM/LLaDA, SUBS: mask never predicted, low-discrepancy mask ratios, no time input)",
                "3:1 bidirectional Gated DeltaNet-2 + bidirectional attention hybrid",
                (
                    "bidirectional GDN-2: shared projections, per-direction short convs and low-rank decay, "
                    "exclusive backward read (Hydra)"
                ),
            ]
        return {
            "architecture": self.config.architecture,
            "layer_types": list(self.config.layer_types),
            "effective_num_layers": self.config.effective_num_layers,
            "shared_layer_indices": [],
            "parameter_report": self.parameter_report(),
            "features": [
                *objective,
                "4-branch Gated Residual stream (gated read, per-branch write scales)",
                "full-rank sigmoid output gates",
                "partial RoPE, per-head QK normalization",
                "SiTU-GLU",
                "tied input-output embeddings",
            ],
        }