# SPDX-License-Identifier: Apache-2.0
"""TT-friendly equivalent forms of the deployed graph, proven on the CPU before device time is spent (the reference
bundles' ``test_fused_host.py`` discipline, REFERENCE_PATTERNS.md 2.12). Each one is exact in real arithmetic (most
are exact in float32 too, as noted) and differs from :mod:`.model` only by float32 summation order;
``tests/test_reference_rewrites_host.py`` checks them against the literal modules with the real weights.

| form | replaces | why (PLAN.md 2.5) |
|---|---|---|
| :func:`pfn_tile_form` | PFN layer 1 on ``concat[h, tile(max h)]`` | ``relu(max_slots(h @ W1a) + max_slots(h) @ W1b + b1)``: no (N, 20, 64) concat, the per-pillar term added after the slot max (exact: ``relu(. + c)`` is monotonic) |
| :func:`pad_pillar_slots` | 20 point slots per pillar | 32 slots (one tile row block per pillar): slots 20-31 repeat slot 0, so both slot maxima are unchanged (max is idempotent; zero rows there would add ``relu(b0)`` to full pillars) |
| :func:`k_split_weights` | ``shared_conv(concat[d0, d1, d2])`` | ``sum_i conv(d_i, W_i)``: the 56 MB concat never exists (PR:P14, S:transfusion:276) |
| :func:`deblock0_as_linear` | Conv 2x2 stride 2 (deblock0) | space-to-depth + linear 256 -> 128 (S:transfusion:271), in case ``conv2d`` loses at 384^2 |
| :func:`deconv_as_linear` | ConvTranspose 2x2 s2 (deblock2) | linear + depth-to-space (the k = 2 device form is ``conv_transpose2d``, PR:P5; this is its oracle) |
| :func:`local_max_masks` | MaxPool p0 + zero border + equality | C22's masked chain ``heat * (m_pass + (heat == maxpool3x3_p1(heat)) * m_eq)`` (exact, PR:P6) |
| :func:`query_pos_from_index` | ``bev_pos`` gather | ``(p % 192 + 0.5, p // 192 + 0.5)`` in float32 (bit-exact; never a bf16 table: S:transfusion:241) |
| class table (``HeadParams.class_table``) | Conv1d of a one-hot | a (5, 128) row gather (bit-exact) |
| :func:`cross_attention_kv_tables` | ``(lidar + pe(bev_pos)) @ Wk + bk`` | ``lidar @ Wk + CK`` with the constant ``CK = pe(bev_pos) @ Wk + bk`` (same for V; S:transfusion:279) |
| :func:`padded_head_attention` | 8 heads x 16 | heads zero-padded to 32 in the projection weights (``ttaw.ops.attention.pad_head_columns`` / ``pad_head_rows``) with ``scale = 1/4`` (SDPA rejects D = 16; exact: PR:P7) |
| :func:`merged_prediction_head_weights` | six Conv1d 128 -> 64 + ReLU -> C heads | one 128 -> 384 + ReLU and a block-diagonal 384 -> 15 (padded to 32) |

numpy weight transforms; torch only inside the ``*_torch`` forwards (imported lazily).
"""
from __future__ import annotations

from typing import Dict, List, Sequence, Tuple

import numpy as np

from ..ttaw.ops.attention import pad_head_columns, pad_head_rows
from .weights import AttentionParams, HeadParams, Layer

__all__ = ["pfn_tile_form", "pad_pillar_slots", "PILLAR_TILE_SLOTS", "PFN_IN_PADDED", "k_split_weights",
           "k_split_torch", "deblock0_as_linear", "space_to_depth_nhwc", "deconv_as_linear", "depth_to_space_nhwc",
           "local_max_masks", "local_max_masked_torch", "query_pos_from_index", "cross_attention_kv_tables",
           "padded_head_attention", "pad_head_dim_rows", "merged_prediction_head_weights", "merged_heads_torch"]

PILLAR_TILE_SLOTS = 32       # one 32-row tile per pillar on the device
PFN_IN_PADDED = 16           # 11 decorated features zero-padded to 16 (32-byte bf16 rows)


def _t(a):
    """float32 torch tensor owning a copy (the ONNX-derived arrays may be read-only)."""
    import torch

    return torch.from_numpy(np.array(a, dtype=np.float32, copy=True))


# --------------------------------------------------------------------------------------------------------- PFN

def pad_pillar_slots(features: np.ndarray, slots: int = PILLAR_TILE_SLOTS, width: int = PFN_IN_PADDED) -> np.ndarray:
    """(P, 20, F) decorated features -> (P, slots, width): features zero-padded to ``width`` columns, slots
    ``20..slots-1`` copies of slot 0 (a real point of every kept pillar)."""
    f = np.asarray(features, dtype=np.float32)
    P, K, F = f.shape
    if slots < K or width < F:
        raise ValueError(f"cannot pad ({K}, {F}) to ({slots}, {width})")
    out = np.zeros((P, slots, width), dtype=np.float32)
    out[:, :K, :F] = f
    out[:, K:, :F] = f[:, :1, :]
    return out


def pfn_tile_form(features, w0, b0, w1a, w1b, b1):
    """The device form of the PFN (torch tensors): ``h = relu(x @ w0 + b0)``, ``m = max_slots h``,
    ``y = relu(max_slots(h @ w1a) + m @ w1b + b1)`` -> (P, 64). ``x`` may carry padded columns (zero rows in ``w0``)
    and repeated slots (:func:`pad_pillar_slots`)."""
    import torch

    h = torch.relu(features @ w0 + b0)                         # (P, S, 32)
    m = h.amax(dim=1)                                          # (P, 32): padded slots included
    z = (h @ w1a).amax(dim=1)                                  # (P, 64): the max taken before the per-pillar term
    return torch.relu(z + m @ w1b + b1)


# ---------------------------------------------------------------------------------------------------- neck / head

def k_split_weights(shared: Layer, split: Sequence[int]) -> List[np.ndarray]:
    """The shared conv weight cut along input channels into the deblocks' channel groups:
    ``conv(concat[d_i], W) == sum_i conv(d_i, W_i)``."""
    if sum(split) != shared.in_channels:
        raise ValueError(f"split {list(split)} does not cover {shared.in_channels} input channels")
    bounds = np.cumsum([0] + list(split))
    return [np.ascontiguousarray(shared.weight[:, a:b]) for a, b in zip(bounds[:-1], bounds[1:])]


def k_split_torch(deblocks, shared: Layer):
    """``sum_i conv(d_i, W_i) + b`` (no activation: the TransFusion shared conv has none), torch NCHW."""
    import torch.nn.functional as F

    parts = k_split_weights(shared, [int(d.shape[1]) for d in deblocks])
    acc = None
    for d, w in zip(deblocks, parts):
        y = F.conv2d(d, _t(w), None, stride=shared.stride, padding=shared.padding)
        acc = y if acc is None else acc + y
    return acc + _t(shared.bias)[None, :, None, None]


def space_to_depth_nhwc(x: np.ndarray, s: int) -> np.ndarray:
    """(H, W, C) -> (H/s, W/s, s*s*C), channel ``(a * s + b) * C + c`` = pixel ``(s*i + a, s*j + b)``."""
    H, W, C = x.shape
    return np.asarray(x).reshape(H // s, s, W // s, s, C).transpose(0, 2, 1, 3, 4).reshape(H // s, W // s, s * s * C)


def deblock0_as_linear(layer: Layer) -> Tuple[np.ndarray, np.ndarray]:
    """Conv k x k stride k (no padding) as ``space_to_depth(x) @ Wl + b``: ``Wl[(a * k + b) * Cin + ci, co] =
    W[co, ci, a, b]``."""
    if layer.kind != "conv" or layer.kernel != layer.stride or layer.padding != 0:
        raise ValueError("deblock0_as_linear needs a conv with kernel == stride and no padding")
    w = layer.weight                                           # (Cout, Cin, k, k)
    cout, cin, k, _ = w.shape
    wl = np.ascontiguousarray(w.transpose(2, 3, 1, 0).reshape(k * k * cin, cout), dtype=np.float32)
    return wl, layer.bias.astype(np.float32)


def deconv_as_linear(layer: Layer) -> Tuple[np.ndarray, np.ndarray]:
    """ConvTranspose with kernel == stride s (no padding) as a linear ``Cin -> s*s*Cout`` on NHWC rows: column
    ``(a * s + b) * Cout + co`` holds ``W[ci, co, a, b]``, bias tiled the same way."""
    if layer.kind != "conv_transpose" or layer.kernel != layer.stride or layer.padding != 0:
        raise ValueError("deconv_as_linear needs a ConvTranspose with kernel == stride and no padding")
    w = layer.weight                                           # (Cin, Cout, s, s)
    cin, cout, s, _ = w.shape
    wl = np.ascontiguousarray(w.transpose(0, 2, 3, 1).reshape(cin, s * s * cout), dtype=np.float32)
    return wl, np.tile(layer.bias, s * s).astype(np.float32)


def depth_to_space_nhwc(y: np.ndarray, H: int, W: int, s: int, cout: int) -> np.ndarray:
    """Rows ``(i, j)`` of ``y`` (H*W, s*s*Cout) -> the (H*s, W*s, Cout) map: ``out[s*i + a, s*j + b, co] =
    y[i*W + j, (a*s + b)*Cout + co]``."""
    return np.asarray(y).reshape(H, W, s, s, cout).transpose(0, 2, 1, 3, 4).reshape(H * s, W * s, cout)


# ----------------------------------------------------------------------------------------------- local max (C22)

def local_max_masks(height: int, width: int, classes: int, kernel: int = 3) -> Tuple[np.ndarray, np.ndarray]:
    """The constant masks of C22's chain for TransFusion's rule (``MaxPool`` k3 s1 **p0** + zero border, every class):
    ``m_eq`` = 1 on interior cells (the pooled window is fully inside), 0 on the ``kernel // 2`` border (always
    suppressed); ``m_pass`` = 0 everywhere (no class bypasses the local max). (classes, H, W) float32."""
    r = kernel // 2
    m_eq = np.zeros((classes, height, width), np.float32)
    m_eq[:, r:height - r, r:width - r] = 1.0
    return np.zeros_like(m_eq), m_eq


def local_max_masked_torch(heat, m_pass, m_eq, kernel: int = 3):
    """``heat * (m_pass + (heat == maxpool_k_s1_p(k//2)(heat)) * m_eq)`` (torch NCHW; the pool's padding is -inf)."""
    import torch.nn.functional as F

    pooled = F.max_pool2d(heat, kernel, 1, kernel // 2)
    return heat * (_t(m_pass)[None] + (heat == pooled).to(heat.dtype) * _t(m_eq)[None])


# ------------------------------------------------------------------------------------------------- query / decoder

def query_pos_from_index(pos: np.ndarray, feat_w: int) -> np.ndarray:
    """(K,) flat BEV cells -> (K, 2) float32 ``(x + 0.5, y + 0.5)`` (exact: integers below 2^24 plus 0.5)."""
    p = np.asarray(pos, dtype=np.int64)
    return np.stack([(p % feat_w).astype(np.float32) + np.float32(0.5),
                     (p // feat_w).astype(np.float32) + np.float32(0.5)], axis=-1)


def cross_attention_kv_tables(attn: AttentionParams, key_pos_embed: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """``CK = pe(bev_pos) @ Wk + bk`` and ``CV = pe(bev_pos) @ Wv + bv`` (36864, 128), float64 then one rounding:
    ``K = lidar^T @ Wk + CK``, ``V = lidar^T @ Wv + CV`` (the 0.6 GMAC position branch leaves the per-frame graph)."""
    kpe = np.asarray(key_pos_embed, dtype=np.float64)
    ck = kpe @ attn.wk.astype(np.float64) + attn.bk.astype(np.float64)
    cv = kpe @ attn.wv.astype(np.float64) + attn.bv.astype(np.float64)
    return ck.astype(np.float32), cv.astype(np.float32)


def pad_head_dim_rows(table: np.ndarray, num_heads: int, head_dim: int, padded: int) -> np.ndarray:
    """(L, H*D) -> (L, H*padded) with each head's columns zero-padded (a constant K / V table in the padded layout)."""
    t = np.asarray(table, dtype=np.float32).reshape(-1, num_heads, head_dim)
    out = np.zeros((t.shape[0], num_heads, padded), np.float32)
    out[..., :head_dim] = t
    return out.reshape(t.shape[0], num_heads * padded)


def padded_head_attention(q_in, k_in, v_in, attn: AttentionParams, *, k_const=None, v_const=None,
                          padded: int = 32):
    """Multi-head attention with the head dim zero-padded to ``padded`` in the weights (torch, (L, 128) inputs):
    ``softmax(scale * Qp Kp^T) Vp`` per head with ``scale = 1 / q_divisor`` (the logical ``1/sqrt(16)``), then the
    row-padded output projection. ``k_const`` / ``v_const`` (L_k, 128): the constant K / V tables
    (:func:`cross_attention_kv_tables`), whose bias they include (then ``k_in @ Wk`` alone is projected)."""
    import torch

    H, D = attn.num_heads, attn.head_dim
    wq, bq = pad_head_columns(attn.wq, attn.bq, H, D, padded)
    wk, bk = pad_head_columns(attn.wk, attn.bk, H, D, padded)
    wv, bv = pad_head_columns(attn.wv, attn.bv, H, D, padded)
    wo = pad_head_rows(attn.wo, H, D, padded)
    q = q_in @ _t(wq) + _t(bq)
    if k_const is None:
        k = k_in @ _t(wk) + _t(bk)
        v = v_in @ _t(wv) + _t(bv)
    else:
        k = k_in @ _t(wk) + _t(pad_head_dim_rows(k_const, H, D, padded))
        v = v_in @ _t(wv) + _t(pad_head_dim_rows(v_const, H, D, padded))
    lq, lk = q.shape[0], k.shape[0]
    qh = q.reshape(lq, H, padded).transpose(0, 1)
    kh = k.reshape(lk, H, padded).transpose(0, 1)
    vh = v.reshape(lk, H, padded).transpose(0, 1)
    o = torch.softmax((qh @ kh.transpose(1, 2)) * attn.scale, dim=-1) @ vh       # (H, Lq, padded)
    return o.transpose(0, 1).reshape(lq, H * padded) @ _t(wo) + _t(attn.bo)


# --------------------------------------------------------------------------------------------- prediction heads

def merged_prediction_head_weights(head: HeadParams, *, pad_to: int = 32
                                   ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[Tuple[str, int, int]]]:
    """The six prediction heads as two matmuls on the (K, 128) decoder rows: ``hidden = relu(x @ Wh + bh)`` with
    ``Wh`` (128, 6*64), and ``out = hidden @ Wo + bo`` with the block-diagonal ``Wo`` (6*64, 15 -> ``pad_to``).
    Returns (Wh, bh, Wo, bo, slices); ``slices`` lists ``(name, start, stop)`` of each head's output channels."""
    wh = np.concatenate([hid.weight[:, :, 0].T for _, hid, _ in head.prediction_heads], axis=1).astype(np.float32)
    bh = np.concatenate([hid.bias for _, hid, _ in head.prediction_heads]).astype(np.float32)
    rows = sum(hid.out_channels for _, hid, _ in head.prediction_heads)
    cols = sum(fin.out_channels for _, _, fin in head.prediction_heads)
    width = -(-cols // pad_to) * pad_to
    wo = np.zeros((rows, width), np.float32)
    bo = np.zeros(width, np.float32)
    slices, r, c = [], 0, 0
    for name, hid, fin in head.prediction_heads:
        wo[r:r + hid.out_channels, c:c + fin.out_channels] = fin.weight[:, :, 0].T
        bo[c:c + fin.out_channels] = fin.bias
        slices.append((name, c, c + fin.out_channels))
        r += hid.out_channels
        c += fin.out_channels
    return wh, bh, wo, bo, slices


def merged_heads_torch(x, merged) -> Dict[str, object]:
    """Apply :func:`merged_prediction_head_weights` to the (K, 128) decoder output -> {name: (1, C, K)}."""
    import torch

    wh, bh, wo, bo, slices = merged
    y = torch.relu(x @ _t(wh) + _t(bh)) @ _t(wo) + _t(bo)                    # (K, 32)
    return {name: y[:, a:b].T[None] for name, a, b in slices}