"""CUDA-graph single-token decoding for Yunmo-Next.

Cached decoding through the regular modules is launch-bound: a token costs
about 5,200 kernel launches across 60 layers while the GPU is busy for roughly
a tenth of the step.  This module replays the whole step as one CUDA graph.

What changes relative to ``YunmoNextForCausalLM.forward_cached``:

- GQA keys/values live in preallocated buffers.  Global layers use a buffer of
  the bucket length, local layers a 2,048-slot ring (keys are stored after
  RoPE, so slot order does not matter); a mask computed from the position
  tensor selects valid slots.
- KDA convolution caches are FLA's own ``[N, D, W]`` buffers, updated in place
  by ``ShortConvolution``; the recurrent state is updated in place by FLA's
  fused recurrent kernel (``inplace_final_state=True``).
- Depth routing is the model's own cached Attention Residual code, unchanged:
  only each layer's mixer call is redirected to the static versions while the
  graph is captured.

Prefill still runs through ``forward_cached``; ``load`` copies its caches into
the static buffers.  One graph is captured per length bucket and reused.

With ``batch > 1`` each row is an independent sequence with its own position
(continuous batching): ``load(cache, row)`` fills a row from a batch-1
prefill, ``release(row)`` parks it, and ``step(tokens)`` advances every
active row.  A 1B model's decode step is bound by reading the weights and by
per-kernel gaps, both nearly independent of the batch, so rows are almost
free until the GPU saturates.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn.functional as F

from .attention import GatedGroupedQueryAttention
from .protocol import CONTROL_BASE

# Per-row greedy summary columns: lexical max value, its index, any-NaN flag,
# then the logits of the 256 protocol IDs.
SUMMARY_COLUMNS = 3 + 256

BUCKETS = (4096, 16384, 65536, 262144)


@dataclass
class _GQABuffers:
    keys: torch.Tensor
    values: torch.Tensor
    is_global: bool


@dataclass
class _KDABuffers:
    state: torch.Tensor
    conv: tuple[torch.Tensor, torch.Tensor, torch.Tensor]


class StaticDecoder:
    """Captured one-token decode step for one length bucket and batch size."""

    def __init__(self, model, length: int, batch: int = 1) -> None:
        if model.training:
            raise RuntimeError("static decoding requires model.eval()")
        self.model = model
        self.length = length
        self.batch = batch
        config = model.config
        self.window = config.attention_window
        param = next(model.parameters())
        device, dtype = param.device, param.dtype
        self.device = device
        self.token = torch.zeros((batch, 1), dtype=torch.long, device=device)
        self.position = torch.zeros((batch, 1), dtype=torch.long, device=device)
        # Parked rows keep rewriting one slot instead of advancing.
        self.active = torch.zeros((batch, 1), dtype=torch.long, device=device)
        self.buffers: dict[int, _GQABuffers | _KDABuffers] = {}
        for index, layer in enumerate(model.layers):
            mixer = layer.mixer
            if isinstance(mixer, GatedGroupedQueryAttention):
                slots = length if mixer.is_global else self.window
                shape = (batch, config.num_key_value_heads, slots, config.attention_head_dim)
                self.buffers[index] = _GQABuffers(
                    torch.zeros(shape, dtype=dtype, device=device),
                    torch.zeros(shape, dtype=dtype, device=device),
                    bool(mixer.is_global),
                )
            else:
                width = mixer.q_conv1d.weight.shape[-1]
                self.buffers[index] = _KDABuffers(
                    torch.zeros((batch, mixer.num_heads, mixer.head_v_dim, mixer.head_k_dim),
                                dtype=torch.float32, device=device),
                    tuple(
                        torch.zeros((batch, conv.weight.shape[0], width), dtype=dtype, device=device)
                        for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d)
                    ),
                )
        self.graph: torch.cuda.CUDAGraph | None = None
        self._host_position = [0] * batch
        self._host_active = [False] * batch
        self.logits: torch.Tensor | None = None
        self.summary: torch.Tensor | None = None

    # ---- per-layer static mixers -------------------------------------------------

    def _gqa(self, layer, buffers: _GQABuffers, hidden: torch.Tensor) -> torch.Tensor:
        mixer = layer.mixer
        normalized = layer.input_norm(hidden)
        raw_q, raw_k, raw_v, gate = mixer._linear_project(normalized)
        query, key, value = mixer._shape_projected(raw_q, raw_k, raw_v, self.position)
        batch, heads, slots, width = buffers.keys.shape
        position = self.position.view(batch, 1)
        slot = position if buffers.is_global else torch.remainder(position, slots)
        index = slot.view(batch, 1, 1, 1).expand(batch, heads, 1, width)
        buffers.keys.scatter_(2, index, key)
        buffers.values.scatter_(2, index, value)
        # Ring slots fill in order from 0, so the first min(position + 1, slots)
        # slots hold exactly the keys this query may see.
        valid = torch.arange(slots, device=hidden.device).view(1, slots) <= position
        # One query token: fold each KV head's query heads (SDPA's GQA maps
        # query head h to KV head h // groups) into the query-length axis and
        # attend with two batched matmuls, reading each KV buffer once.  SDPA
        # with ``enable_gqa`` expands KV to every query head (~2 ms per layer
        # at batch 16) and SDPA on the folded shape picks a slow FP32 kernel;
        # measured per layer at 4,096 slots: 0.17-0.32 ms here vs 0.27-3.9 ms.
        # Same scale as the model (SDPA default 1/sqrt(head_dim)).
        query_heads = query.shape[1]
        grouped = query.reshape(batch, heads, query_heads // heads, width)
        scores = torch.matmul(grouped, buffers.keys.transpose(-1, -2)) * width**-0.5
        scores = scores.masked_fill(~valid.view(batch, 1, 1, slots), float("-inf"))
        output = torch.matmul(torch.softmax(scores, dim=-1), buffers.values)
        output = output.reshape(batch, query_heads, 1, width)
        return mixer._apply_output_gate(output, gate)

    def _kda(self, layer, buffers: _KDABuffers, hidden: torch.Tensor) -> torch.Tensor:
        from einops import rearrange
        from fla.ops.kda.fused_recurrent import fused_recurrent_kda_fwd

        mixer = layer.mixer
        x = layer.input_norm(hidden)
        conv_q, conv_k, conv_v = buffers.conv
        q, _ = mixer.q_conv1d(x=mixer.q_proj(x), cache=conv_q, output_final_state=True)
        k, _ = mixer.k_conv1d(x=mixer.k_proj(x), cache=conv_k, output_final_state=True)
        v, _ = mixer.v_conv1d(x=mixer.v_proj(x), cache=conv_v, output_final_state=True)
        g = rearrange(mixer.f_proj(x), "... (h d) -> ... h d", d=mixer.head_k_dim)
        beta = mixer.b_proj(x)
        q, k = (rearrange(t, "... (h d) -> ... h d", d=mixer.head_k_dim) for t in (q, k))
        v = rearrange(v, "... (h d) -> ... h d", d=mixer.head_v_dim)
        output, _ = fused_recurrent_kda_fwd(
            q=q, k=k, v=v, g=g, beta=beta, A_log=mixer.A_log, dt_bias=mixer.dt_bias,
            scale=k.shape[-1] ** -0.5, initial_state=buffers.state, inplace_final_state=True,
            output_final_state=True, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True,
            use_beta_sigmoid_in_kernel=True, allow_neg_eigval=mixer.allow_neg_eigval,
            lower_bound=mixer.lower_bound, cu_seqlens=None, state_v_first=True,
        )
        gate = rearrange(mixer.g_proj(x), "... (h d) -> ... h d", d=mixer.head_v_dim)
        output = rearrange(mixer.o_norm(output, gate), "b t h d -> b t (h d)")
        return mixer.o_proj(output)

    # ---- step and capture --------------------------------------------------------

    def _step(self) -> torch.Tensor:
        model = self.model
        originals = {}
        for index, layer in enumerate(model.layers):
            buffers = self.buffers[index]
            static = self._gqa if isinstance(buffers, _GQABuffers) else self._kda

            def mixer_output_cached(hidden, position_ids, *, fla_cache, attention_cache,
                                    _layer=layer, _buffers=buffers, _static=static):
                return _static(_layer, _buffers, hidden), None

            originals[index] = layer.__dict__.get("mixer_output_cached")
            layer.mixer_output_cached = mixer_output_cached
        try:
            state = _NullState()
            hidden = model._forward_hidden_two_phase_attention_residual_cached(
                model.embed_tokens(self.token), self.position, state, None, None,
            )
            return F.linear(hidden, model.embed_tokens.weight)[:, -1]
        finally:
            for index, layer in enumerate(model.layers):
                if originals[index] is None:
                    del layer.mixer_output_cached
                else:
                    layer.mixer_output_cached = originals[index]

    def _tensors(self):
        for buffers in self.buffers.values():
            if isinstance(buffers, _GQABuffers):
                yield buffers.keys
                yield buffers.values
            else:
                yield buffers.state
                yield from buffers.conv

    @torch.inference_mode()
    def capture(self, warmup: int = 3) -> None:
        stream = torch.cuda.Stream()
        stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(stream):
            for _ in range(warmup):
                _greedy_summary(self._step())
        torch.cuda.current_stream().wait_stream(stream)
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            self.logits = self._step()
            self.summary = _greedy_summary(self.logits)
        # Warm-up and capture ran real steps on the buffers.  ``load`` rewrites
        # every buffer of the row it fills, so clearing (instead of keeping a
        # snapshot copy, which doubled peak memory) is enough.
        for tensor in self._tensors():
            tensor.zero_()
        self.position.zero_()
        self.active.zero_()
        self.graph = graph

    # ---- runtime -----------------------------------------------------------------

    @torch.inference_mode()
    def load(self, cache, row: int = 0) -> int:
        """Copy a batch-1 ``forward_cached`` prefill cache into ``row``."""

        position = int(cache.seen_tokens)
        if position >= self.length:
            raise ValueError("prompt does not fit this decode bucket")
        for index, buffers in self.buffers.items():
            if isinstance(buffers, _GQABuffers):
                buffers.keys[row].zero_()
                buffers.values[row].zero_()
                entry = cache.attention.get(index)
                if entry is None:
                    continue
                count = entry.keys.shape[2]
                positions = torch.arange(position - count, position, device=self.device)
                slots = positions if buffers.is_global else torch.remainder(positions, self.window)
                buffers.keys[row].index_copy_(1, slots, entry.keys[0])
                buffers.values[row].index_copy_(1, slots, entry.values[0])
            else:
                layer_state = cache.fla_cache[index]
                buffers.state[row].copy_(layer_state["recurrent_state"][0])
                for target, source in zip(buffers.conv, layer_state["conv_state"], strict=True):
                    target[row].copy_(source[0])
        self.position[row].fill_(position)
        self.active[row].fill_(1)
        self._host_position[row] = position
        self._host_active[row] = True
        return position

    def release(self, row: int) -> None:
        """Park ``row``; its outputs are ignored until the next ``load``."""

        self.active[row].fill_(0)
        self.position[row].fill_(0)
        self._host_active[row] = False

    @torch.inference_mode()
    def step(self, tokens) -> torch.Tensor:
        """``advance`` and return a copy of the ``[batch, vocab]`` logits."""

        self.advance(tokens)
        return self.logits.clone()

    @torch.inference_mode()
    def advance(self, tokens) -> None:
        """Advance every active row by one token.

        ``tokens`` is an int (batch 1) or one int per row; parked rows' entries
        are ignored.  Afterwards ``logits`` and ``summary`` hold the new rows
        until the next call.
        """

        if self.graph is None:
            raise RuntimeError("capture() must run before step()")
        if isinstance(tokens, int):
            tokens = [tokens]
        if len(tokens) != self.batch:
            raise ValueError("step needs one token per row")
        # A global-layer write past the buffer would be an out-of-bounds
        # scatter inside the graph; refuse on the host instead.
        for row in range(self.batch):
            if self._host_active[row]:
                if self._host_position[row] >= self.length:
                    raise RuntimeError("decode exceeded the static decoder's length bucket")
                self._host_position[row] += 1
        self.token.copy_(torch.tensor(tokens, dtype=torch.long).view(self.batch, 1), non_blocking=False)
        self.graph.replay()
        self.position.add_(self.active)


def _greedy_summary(logits: torch.Tensor) -> torch.Tensor:
    # Float32 throughout: the index (< 2**24) is exact in float32 but not in
    # float16, and float16 logits widen exactly.
    lexical = logits[:, :CONTROL_BASE]
    index = lexical.argmax(dim=-1, keepdim=True)
    return torch.cat(
        (
            lexical.gather(-1, index).float(),
            index.float(),
            torch.isnan(logits).any(dim=-1, keepdim=True).float(),
            logits[:, CONTROL_BASE:].float(),
        ),
        dim=-1,
    )


class _NullState:
    """Stands in for ``YunmoNextCUDAState``: static mixers ignore it."""

    fla_cache = None
    attention: dict = {}


def bucket_for(length: int) -> int:
    for bucket in BUCKETS:
        if length <= bucket:
            return bucket
    raise ValueError("requested length exceeds the largest decode bucket")