"""Layer-addressable decoder with a static per-layer KV cache.

The wrapper drives the submodules of a Hugging Face dense decoder (Qwen3, Llama)
one layer range at a time. Every layer owns a KV buffer of size `max_len`; a
token's K/V are written at its absolute position, so layers may hold different
sets of positions while drafting (early exits) and verification can overwrite
approximate entries in place.

Attention for a chunk of `n` new rows at positions [s, s+n) reads K/V
positions [0, s+n) of the same layer, with a causal mask inside the chunk.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn.functional as F


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


@dataclass
class ModelSpec:
    num_layers: int
    hidden_size: int
    vocab_size: int
    num_heads: int
    num_kv_heads: int
    head_dim: int
    has_qk_norm: bool


def use_row_invariant_blas() -> None:
    """Select rocBLAS on ROCm: row-invariant for 1..64 rows on MI210 (job 893629),
    unlike hipBLASLt (o_proj breaks at 17 rows, q_proj at 24). No-op elsewhere."""
    if torch.cuda.is_available() and torch.version.hip:
        torch.backends.cuda.preferred_blas_library("cublas")  # maps to rocBLAS on ROCm


class LayeredDecoder:
    """Run a frozen HF causal LM layer by layer with a positional KV cache."""

    def __init__(self, hf_model, max_len: int, canonical_attention: bool = True,
                 max_rows: int = 16):
        self.hf = hf_model
        # Canonical attention: every query row uses the one-query SDPA call over its
        # exact prefix, as in one-token AR. On MI210/bf16 the batched masked call
        # is not row-invariant while all projections and norms are (job 893599).
        self.canonical_attention = canonical_attention
        # GEMMs were verified row-invariant up to 16 rows (job 893599); larger
        # chunks are processed in causal sub-chunks of at most `max_rows` rows.
        self.max_rows = max_rows
        base = hf_model.model
        cfg = hf_model.config
        self.embed_tokens = base.embed_tokens
        self.layers = base.layers
        self.norm = base.norm
        self.lm_head = hf_model.lm_head
        self.rotary_emb = base.rotary_emb
        attn0 = self.layers[0].self_attn
        self.spec = ModelSpec(
            num_layers=len(self.layers),
            hidden_size=cfg.hidden_size,
            vocab_size=cfg.vocab_size,
            num_heads=cfg.num_attention_heads,
            num_kv_heads=cfg.num_key_value_heads,
            head_dim=attn0.head_dim,
            has_qk_norm=hasattr(attn0, "q_norm"),
        )
        self.scaling = attn0.scaling
        self.max_len = max_len
        p = next(hf_model.parameters())
        self.device, self.dtype = p.device, p.dtype
        shape = (self.spec.num_layers, self.spec.num_kv_heads, max_len, self.spec.head_dim)
        self.k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
        self.v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
        # Rotary tables for every position, computed once.
        pos = torch.arange(max_len, device=self.device).unsqueeze(0)
        dummy = torch.zeros(1, 1, self.spec.head_dim, device=self.device, dtype=self.dtype)
        cos, sin = self.rotary_emb(dummy, pos)
        self.cos, self.sin = cos[0], sin[0]  # (max_len, head_dim)
        self.counters = {"layer_rows": 0, "kv_fill_rows": 0, "vocab_proj_rows": 0}

    # ------------------------------------------------------------------ helpers
    def reset_counters(self) -> None:
        for k in self.counters:
            self.counters[k] = 0

    def embed(self, ids: torch.Tensor) -> torch.Tensor:
        """ids: (n,) -> (n, d)"""
        return self.embed_tokens(ids)

    def _rope(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
        # x: (heads, n, head_dim); pos: (n,) absolute positions
        cos = self.cos.index_select(0, pos).unsqueeze(0)
        sin = self.sin.index_select(0, pos).unsqueeze(0)
        return x * cos + _rotate_half(x) * sin

    def _kv(self, attn, x_normed: torch.Tensor, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        n = x_normed.shape[0]
        hd = self.spec.head_dim
        k = attn.k_proj(x_normed).view(n, -1, hd)
        if self.spec.has_qk_norm:
            k = attn.k_norm(k)
        k = self._rope(k.transpose(0, 1), pos)
        v = attn.v_proj(x_normed).view(n, -1, hd).transpose(0, 1)
        return k, v

    def _positions(self, start: int, n: int, positions) -> tuple[torch.Tensor, list[int]]:
        if positions is None:
            plist = list(range(start, start + n))
        else:
            plist = [int(p) for p in positions]
            assert len(plist) == n and all(a < b for a, b in zip(plist, plist[1:])), "ascending positions"
        return torch.tensor(plist, device=self.device), plist

    # ------------------------------------------------------------------ layers
    def run_layers(self, h: torch.Tensor, start: int, first: int, last: int,
                   positions=None) -> torch.Tensor:
        """Apply layers [first, last) to hidden rows h (n, d).

        Rows sit at positions [start, start+n), or at the given ascending
        `positions` (rows of a draft block caught up to a deeper layer). K/V of
        every processed layer are written at those positions; row i attends to
        positions [0, pos_i] of the same layer.
        """
        n = h.shape[0]
        pos, plist = self._positions(start, n, positions)
        if self.canonical_attention and n > self.max_rows:
            m = self.max_rows  # causal: earlier rows never depend on later rows
            return torch.cat([self.run_layers(h[i:i + m], 0, first, last, plist[i:i + m])
                              for i in range(0, n, m)])
        hd = self.spec.head_dim
        end = plist[-1] + 1
        mask = None
        if n > 1 and not self.canonical_attention:
            mask = torch.arange(0, end, device=h.device).unsqueeze(0) <= pos.unsqueeze(1)  # (n, end)
        for li in range(first, last):
            layer = self.layers[li]
            attn = layer.self_attn
            x = layer.input_layernorm(h)
            q = attn.q_proj(x).view(n, -1, hd)
            if self.spec.has_qk_norm:
                q = attn.q_norm(q)
            q = self._rope(q.transpose(0, 1), pos)
            k, v = self._kv(attn, x, pos)
            self.k_cache[li].index_copy_(1, pos, k)
            self.v_cache[li].index_copy_(1, pos, v)
            if n == 1 or not self.canonical_attention:
                o = F.scaled_dot_product_attention(
                    q.unsqueeze(0),
                    self.k_cache[li, :, :end].unsqueeze(0),
                    self.v_cache[li, :, :end].unsqueeze(0),
                    attn_mask=mask,
                    scale=self.scaling,
                    enable_gqa=True,
                )[0]
            else:
                o = torch.cat([F.scaled_dot_product_attention(
                    q[:, i:i + 1].unsqueeze(0),
                    self.k_cache[li, :, :p + 1].unsqueeze(0),
                    self.v_cache[li, :, :p + 1].unsqueeze(0),
                    scale=self.scaling,
                    enable_gqa=True,
                )[0] for i, p in enumerate(plist)], dim=1)
            o = attn.o_proj(o.transpose(0, 1).reshape(n, -1))
            h = h + o
            h = h + layer.mlp(layer.post_attention_layernorm(h))
        self.counters["layer_rows"] += n * (last - first)
        return h

    def fill_kv_from_hidden(self, h: torch.Tensor, start: int, first: int, last: int) -> None:
        """CALM-style state copy (regime B): K/V of layers [first, last) computed from h itself.

        Used for tokens that exited early, so that deeper draft tokens can attend
        to them. Only affects draft quality, never verified outputs.
        """
        pos, _ = self._positions(start, h.shape[0], None)
        for li in range(first, last):
            layer = self.layers[li]
            k, v = self._kv(layer.self_attn, layer.input_layernorm(h), pos)
            self.k_cache[li].index_copy_(1, pos, k)
            self.v_cache[li].index_copy_(1, pos, v)
        self.counters["kv_fill_rows"] += h.shape[0] * max(0, last - first)

    def final_logits(self, h: torch.Tensor) -> torch.Tensor:
        self.counters["vocab_proj_rows"] += h.shape[0]
        return self.lm_head(self.norm(h)).float()

    # ------------------------------------------------------------------ full passes
    def prefill(self, ids: torch.Tensor) -> torch.Tensor:
        """Process the whole prompt; returns last-position logits (V,).

        The prompt is shared by every decoding method, so it uses the fast batched
        attention (row invariance is only needed for generated rows).
        """
        canon = self.canonical_attention
        self.canonical_attention = False
        try:
            h = self.run_layers(self.embed(ids), 0, 0, self.spec.num_layers)
        finally:
            self.canonical_attention = canon
        return self.final_logits(h[-1:])[0]

    def full_step(self, ids: torch.Tensor, start: int) -> torch.Tensor:
        """All layers on rows ids at positions [start, ...); returns logits (n, V)."""
        h = self.run_layers(self.embed(ids), start, 0, self.spec.num_layers)
        return self.final_logits(h)