# SPDX-License-Identifier: Apache-2.0
"""fp32 torch rebuild of the deployed PTv3 v4.0 graphs (encoder, seg head, det head) on the host tables of
``tt_ptv3.host.preprocess`` (identical to what the device consumes), with per-module taps.

Semantics reproduced from the graphs (asserted in ``weights.py``; research/ptv3/SPEC.md 4):

- GELU is the exact erf form; LayerNorm eps 1e-5; BN folded where ``weights.py`` says;
- sub-manifold conv = ``sum_k x[nmap[:, k]] @ W_k`` over the 27 taps of the host neighbour map (missing taps = 0);
- serialized attention: block ``b`` uses order ``b % 2``; ``Npad = ceil(N / 512) * 512``, the sequence row ``j``
  is voxel ``order[min(j, N - 1)]`` (pad tokens are copies of the last voxel in curve order, no mask), heads of 16,
  ``softmax(0.25 * q k^T) v`` per 512-row patch, rows back through ``inverse``;
- pooling: ``proj`` Linear -> segment max over the CSR children -> BN affine -> GELU;
- seg decoder: ``x = GELU(skip(pf_s)) + GELU(up(x))[cluster_s]``; at stages 3 / 2 one block whose CPE reads
  ``GELU(skip(pf_s))`` (the export quirk), order 0; classifier -> softmax -> first-max argmax;
- det head: fusion, ReLU point projection, BEV scatter-max into a zero 256 x 256 canvas as a gather + max over the
  host's per-cell candidate lists (``host.preprocess.bev_index``; exact: max of non-negative values with a zero
  row), 12 dense 3x3 convs, heat logits, sigmoid, 3x3 local max on the interior of classes 0-3 with the heat itself
  as the seed (border cells and classes 4-6 pass), ``heat * (heat == local_max)``, ONNX-semantics TopK 500 over the
  class-major flattening (ties: lower index first; ``topk_indices=`` teacher-forces it), query features + class
  encoding, ``query_pos = bev_pos[pos]``, one post-norm decoder layer (self-attention with Q = K = V = query + pos
  MLP; cross-attention to the 65,536 BEV keys + key pos MLP), FFN, six prediction heads (``center += query_pos``).

Taps: pass ``taps={}`` (a dict, filled with float32 numpy arrays) and optionally ``tap_filter`` (a callable on the
name); names are listed in :data:`TAP_NAMES` (``{s}`` / ``{b}`` / ``{i}`` expanded).
"""
from __future__ import annotations

from typing import Callable, Dict, List, Optional, Sequence

import numpy as np

from .config import HEAD_DIM, PATCH_SIZE
from .weights import Block, Conv2d, DetHead, Encoder, Linear, MHA, Norm, SegHead, SubmConv

__all__ = ["TAP_NAMES", "encoder_forward", "seg_head_forward", "det_head_forward", "topk_onnx", "ROW_CHUNK",
           "PATCH_CHUNK"]

TAP_NAMES = (
    "embed", "enc{s}.down", "enc{s}.block{b}.cpe", "enc{s}.block{b}.attn", "enc{s}.block{b}", "point_feat_{s}",
    "dec{s}.up", "dec{s}.skip", "dec{s}.sum", "dec{s}.block0", "seg.logits", "seg.probs", "seg.labels",
    "det.fusion", "det.point_proj", "det.bev", "det.bev_refine", "det.stem", "det.block{i}", "det.shared",
    "det.dense_heatmap", "det.heat", "det.heat_nms", "det.topk_indices", "det.query_feat", "det.query_pos",
    "det.query_pos_embed", "det.key_pos_embed", "det.self", "det.norm1", "det.cross", "det.norm2", "det.ffn",
    "det.norm3", "det.head.{name}")
ROW_CHUNK = 32768        # rows per sparse-conv gather chunk (host memory only; results do not depend on it)
PATCH_CHUNK = 32         # patches per attention chunk


def _torch():
    import torch
    return torch


class _Taps:
    def __init__(self, taps: Optional[Dict[str, np.ndarray]], tap_filter: Optional[Callable[[str], bool]]):
        self.taps, self.filter = taps, tap_filter

    def __call__(self, name: str, x):
        if self.taps is not None and (self.filter is None or self.filter(name)):
            self.taps[name] = x.detach().cpu().numpy().astype(np.float32, copy=True) if hasattr(x, "detach") \
                else np.asarray(x)
        return x


def _t(a):
    torch = _torch()
    a = np.ascontiguousarray(a)
    return torch.from_numpy(a if a.flags.writeable else a.copy())


def linear(x, p: Linear):
    torch = _torch()
    return torch.nn.functional.linear(x, _t(p.w), None if p.b is None else _t(p.b))


def layer_norm(x, n: Norm):
    torch = _torch()
    return torch.nn.functional.layer_norm(x, (x.shape[-1],), _t(n.g), _t(n.b), n.eps)


def gelu(x):
    return _torch().nn.functional.gelu(x, approximate="none")


def subm_conv(x, nmap: np.ndarray, conv: SubmConv):
    """Gather-form sub-manifold conv (zero row for missing taps), float32."""
    torch = _torch()
    n, cin = x.shape
    cout = conv.w.shape[0]
    xz = torch.cat([x, torch.zeros(1, cin, dtype=x.dtype)], 0)
    idx = torch.from_numpy(np.where(nmap >= 0, nmap, n).astype(np.int64))
    wmat = _t(conv.w).permute(1, 2, 0).reshape(27 * cin, cout)        # tap-major rows k * Cin + j
    out = torch.empty(n, cout, dtype=x.dtype)
    for r0 in range(0, n, ROW_CHUNK):
        g = xz[idx[r0:r0 + ROW_CHUNK]].reshape(-1, 27 * cin)
        out[r0:r0 + ROW_CHUNK] = g @ wmat
    return out + _t(conv.b)


def serialized_attention(x, blk: Block, order: np.ndarray, inverse: np.ndarray):
    """LN'd input ``x`` -> attention output before ``proj``'s residual (``proj`` applied)."""
    torch = _torch()
    n, c = x.shape
    h = blk.heads
    npad = -(-n // PATCH_SIZE) * PATCH_SIZE
    seq = torch.from_numpy(np.asarray(order, np.int64)[np.minimum(np.arange(npad), n - 1)])
    qkv = linear(x, blk.qkv)[seq].reshape(-1, PATCH_SIZE, 3, h, HEAD_DIM).permute(2, 0, 3, 1, 4)
    q, k, v = qkv[0], qkv[1], qkv[2]                                  # [P, H, 512, 16]
    out = torch.empty(q.shape[0], PATCH_SIZE, h, HEAD_DIM, dtype=x.dtype)
    for p0 in range(0, q.shape[0], PATCH_CHUNK):
        sl = slice(p0, p0 + PATCH_CHUNK)
        a = torch.softmax((q[sl] * (HEAD_DIM ** -0.5)) @ k[sl].transpose(-2, -1), dim=-1)
        out[sl] = (a @ v[sl]).transpose(1, 2)
    o = out.reshape(-1, c)[torch.from_numpy(np.asarray(inverse, np.int64))]
    return linear(o, blk.proj)


def block_forward(x, blk: Block, nmap, order, inverse, tap: _Taps, cpe_in=None, name: str = ""):
    cpe_in = x if cpe_in is None else cpe_in
    x = x + layer_norm(linear(subm_conv(cpe_in, nmap, blk.cpe_conv), blk.cpe_linear), blk.cpe_norm)
    tap(f"{name}.cpe", x)
    x = x + serialized_attention(layer_norm(x, blk.norm1), blk, order, inverse)
    tap(f"{name}.attn", x)
    x = x + linear(gelu(linear(layer_norm(x, blk.norm2), blk.fc1)), blk.fc2)
    return tap(name, x)


def segment_max(x, indptr: np.ndarray):
    torch = _torch()
    indptr = np.asarray(indptr, np.int64)
    lens = torch.from_numpy(indptr[1:] - indptr[:-1])
    seg = torch.repeat_interleave(torch.arange(lens.shape[0]), lens)
    out = torch.full((lens.shape[0], x.shape[1]), -float("inf"), dtype=x.dtype)
    return out.scatter_reduce(0, seg[:, None].expand(-1, x.shape[1]), x[:int(indptr[-1])], reduce="amax",
                              include_self=True)


def encoder_forward(enc: Encoder, frame, *, taps: Optional[Dict[str, np.ndarray]] = None,
                    tap_filter: Optional[Callable[[str], bool]] = None) -> List:
    """``frame`` (``host.preprocess.Frame``) -> ``[point_feat_0..4]`` (torch float32)."""
    torch = _torch()
    tap = _Taps(taps, tap_filter)
    x = tap("embed", gelu(linear(torch.from_numpy(frame.voxels.feat), enc.embed)))
    outs = []
    for s, st in enumerate(enc.stages):
        lv = frame.levels[s]
        if s > 0:
            pooled = segment_max(linear(x, st.down), frame.pools[s - 1].indptr)
            x = tap(f"enc{s}.down", gelu(pooled * _t(st.down_norm.scale) + _t(st.down_norm.shift)))
        for b, blk in enumerate(st.blocks):
            x = block_forward(x, blk, lv.nmap, lv.order[blk.order], lv.inverse[blk.order], tap, name=f"enc{s}.block{b}")
        outs.append(tap(f"point_feat_{s}", x))
    return outs


def seg_head_forward(seg: SegHead, point_feats: Sequence, frame, *, taps=None, tap_filter=None):
    """-> ``(labels int64 [N0], probs float32 [N0, 15])`` (torch)."""
    torch = _torch()
    tap = _Taps(taps, tap_filter)
    x = point_feats[4]
    for s in (3, 2, 1, 0):
        d = seg.dec[s]
        up = tap(f"dec{s}.up", gelu(linear(x, d.up)))
        skip = tap(f"dec{s}.skip", gelu(linear(point_feats[s], d.skip)))
        x = tap(f"dec{s}.sum", skip + up[torch.from_numpy(frame.pools[s].cluster)])
        if d.block is not None:
            lv = frame.levels[s]
            x = block_forward(x, d.block, lv.nmap, lv.order[0], lv.inverse[0], tap, cpe_in=skip, name=f"dec{s}.block0")
    logits = tap("seg.logits", linear(x, seg.classifier))
    probs = tap("seg.probs", torch.softmax(logits, dim=1))
    labels = probs.argmax(1)
    tap("seg.labels", labels)
    return labels, probs


# ---------------------------------------------------------------------------------------------- det head
def conv2d(x, c: Conv2d, residual=None):
    torch = _torch()
    y = torch.nn.functional.conv2d(x, _t(c.w), _t(c.b), padding=c.padding, dilation=c.dilation)
    if residual is not None:
        y = y + residual
    return torch.relu(y) if c.relu else y


def topk_onnx(scores_flat: np.ndarray, k: int) -> np.ndarray:
    """ONNX ``TopK(largest=1, sorted=1)`` indices with ties broken by the lower index (ORT's order)."""
    s = np.asarray(scores_flat, np.float32).reshape(-1)
    return np.argsort(-s, kind="stable")[:k].astype(np.int64)


def mha(q_in, kv_in, m: MHA, v_in=None):
    """``[Sq, D]`` x ``[Sk, D]`` -> ``[Sq, D]``: projections, heads, ``softmax((q / div) k^T) v``, out projection."""
    torch = _torch()
    v_in = kv_in if v_in is None else v_in
    hd = m.q.w.shape[0] // m.heads
    q = linear(q_in, m.q).reshape(-1, m.heads, hd).transpose(0, 1) / m.scale_div
    k = linear(kv_in, m.k).reshape(-1, m.heads, hd).transpose(0, 1)
    v = linear(v_in, m.v).reshape(-1, m.heads, hd).transpose(0, 1)
    o = torch.empty_like(q)
    for h in range(m.heads):                         # per head: the 500 x 65,536 cross scores stay ~130 MB
        o[h] = torch.softmax(q[h] @ k[h].T, dim=-1) @ v[h]
    return linear(o.transpose(0, 1).reshape(q_in.shape[0], -1), m.out)


def det_head_forward(det: DetHead, point_feats: Sequence, frame, bev_index: np.ndarray, *,
                     topk_indices: Optional[np.ndarray] = None, taps=None, tap_filter=None) -> Dict[str, np.ndarray]:
    """-> the 9 outputs of ``ptv3_det3d_head.onnx`` (numpy, ONNX shapes and dtypes). ``bev_index``: the host's
    ``(65536, 9)`` candidate table whose sentinel is the row count of stage 3 (``host.preprocess.bev_index``)."""
    torch = _torch()
    tap = _Taps(taps, tap_filter)
    pf3, pf4 = point_feats[3], point_feats[4]
    n3 = pf3.shape[0]
    cluster3 = torch.from_numpy(frame.pools[3].cluster)
    fused = tap("det.fusion", gelu(linear(pf3, det.fusion_skip)) + gelu(linear(pf4, det.fusion_proj))[cluster3])
    proj = tap("det.point_proj", torch.relu(linear(fused, det.point_proj)))
    table = torch.cat([proj, torch.zeros(1, proj.shape[1], dtype=proj.dtype)], 0)
    idx = torch.from_numpy(np.where(np.asarray(bev_index) >= n3, n3, bev_index).astype(np.int64))
    bev_rows = table[idx].amax(1)                                      # [65536, 128] (empty cells: 0)
    tap("det.bev", bev_rows)
    size = int(round(bev_rows.shape[0] ** 0.5))
    x = bev_rows.reshape(1, size, size, -1).permute(0, 3, 1, 2).contiguous()
    x = tap("det.bev_refine", conv2d(x, det.bev_refine))
    x = tap("det.stem", conv2d(x, det.stem))
    for i, (c1, c2) in enumerate(det.blocks):
        x = tap(f"det.block{i}", conv2d(conv2d(x, c1), c2, residual=x))
    feat = tap("det.shared", conv2d(x, det.shared))                   # [1, 128, H, W]
    logits = tap("det.dense_heatmap", conv2d(conv2d(feat, det.heat_hidden), det.heat_out))
    heat = tap("det.heat", torch.sigmoid(logits))
    local = heat.clone()
    pc = list(det.pooled_classes)
    local[:, pc, 1:-1, 1:-1] = torch.nn.functional.max_pool2d(heat[:, pc], 3, stride=1, padding=0)
    heat_nms = tap("det.heat_nms", heat * (heat == local).to(heat.dtype))
    C = heat.shape[1]
    cells = size * size
    flat = heat_nms.reshape(1, C, cells)
    if topk_indices is None:
        topk_indices = topk_onnx(flat.reshape(-1).numpy(), det.num_proposals)
    top = np.asarray(topk_indices, np.int64).reshape(-1)
    tap("det.topk_indices", top)
    labels = top // cells
    pos = top % cells
    pos_t = torch.from_numpy(pos)
    feat_flat = feat.reshape(1, feat.shape[1], cells)
    onehot = torch.nn.functional.one_hot(torch.from_numpy(labels), C).to(feat.dtype)          # [K, C]
    query = feat_flat[0][:, pos_t].T + linear(onehot, det.class_encoding)                       # [K, 128]
    tap("det.query_feat", query)
    qpos = torch.from_numpy(det.bev_pos)[pos_t]                                                # [K, 2]
    tap("det.query_pos", qpos)
    qpe = tap("det.query_pos_embed", linear(torch.relu(linear(qpos, det.query_pos[0])), det.query_pos[1]))
    kpe = tap("det.key_pos_embed", linear(torch.relu(linear(torch.from_numpy(det.bev_pos), det.key_pos[0])),
                                          det.key_pos[1]))
    q = query + qpe
    q = tap("det.self", q + mha(q, q, det.self_attn))
    q = tap("det.norm1", layer_norm(q, det.norm1))
    keys = feat_flat[0].T + kpe                                                                 # [65536, 128]
    q = tap("det.cross", q + mha(q, keys, det.cross_attn))
    q = tap("det.norm2", layer_norm(q, det.norm2))
    q = tap("det.ffn", q + linear(torch.relu(linear(q, det.ffn[0])), det.ffn[1]))
    q = tap("det.norm3", layer_norm(q, det.norm3))
    out: Dict[str, np.ndarray] = {}
    for name, (h0, h1) in det.heads.items():
        y = linear(torch.relu(linear(q, h0)), h1)                                               # [K, c]
        if name == "center":
            y = y + qpos
        tap(f"det.head.{name}", y)
        out[name] = y.T.unsqueeze(0).numpy().astype(np.float32)
    out["dense_heatmap"] = logits.numpy().astype(np.float32)
    out["query_heatmap_score"] = flat[:, :, pos_t].numpy().astype(np.float32)
    out["query_labels"] = labels.reshape(1, -1).astype(np.int64)
    return out