# SPDX-License-Identifier: Apache-2.0
"""The TT device dataflow of FRNet (PORT_LOG.md "trace layout"), run on the CPU in fp32 and proven equal to the
literal graph (``reference/model.py``) by ``tests/test_reference_rewrites_host.py``. Each rewrite is exact in real
arithmetic; in float32 they differ only by summation order:

- **points sorted by cell and padded** to ``N_cap`` (``host/preprocess.py`` ``FramePlan``): the network is
  permutation-equivariant in the points; padded rows form one trailing sentinel segment whose results nothing reads;
- **host features**: the 8-d decoration and ``pre_norm`` in float64, one rounding (``decorate_points``);
- **segmented max** (6 sites: ``ScatterElements(max)``) as the exact **log-step** form of PLAN.md 1.3 (the K1
  kernel's stock-op fallback): round ``r`` takes ``x = max(x, x[idx_r])`` with the host shift table ``idx_r``,
  then one gather of each segment's last row (``seg_last``) -> ``(M_cap, C)``; :func:`segmax_logstep`;
- **frustum2pixel** (6 ``ScatterND``) as a gather: ``image[pixel] = table[pix2vox_s[pixel]]`` from the voxel table
  with one appended **zero row** (index ``M_cap``) for empty pixels -- ``ttaw.ops.gather.scatter_rows`` semantics
  (C19); duplicates at s > 1 resolved on the host (last voxel wins);
- **pixel2point** as a gather of NHWC rows (``point2pix_s``);
- **concat-free convs**: ``conv(concat(a, b))`` = ``conv_a(a) + conv_b(b)`` for the fusion convs (256 = 128 + 128)
  and the 5-way K-split ``fuse_layer1`` (640 = 5 x 128, probe P14: 8.3 vs 14.2 ms);
- **concat-free point layers**: ``linear(concat(a, b))`` = ``a @ W_a + b @ W_b + bias`` (point stem, point fusion,
  ``point_fuse_layer1``), which also lets the two halves keep different dtypes;
- **align-corners bilinear** as interpolation matrices ``A_h X A_w^T`` (C18 ``interp_matrix``, float64 -> float32).

:class:`Emulation` optionally rounds layer operands the way the device stores them (bf16, or the TF32-like
mantissa of the device's fp32 matmul, probe P12) -- a first-order CPU model to size the fp32 islands of the
port before the device runs (``code/scripts/precision_study.py``); it models no math-fidelity phases or tile
layouts, so its numbers are indicative only.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Sequence

import numpy as np
import torch
import torch.nn.functional as F

from ..host.preprocess import FramePlan, segmax_rounds, segmax_shift_tables
from ..ttaw.golden import NULL_TAPS, TapRegistry
from ..ttaw.ops.upsample import interp_matrix
from .model import TorchLayers, _t
from .weights import Conv, FRNetWeights, Linear

__all__ = ["Emulation", "round_mode", "segmax_logstep", "FRNetPlanTorch", "plan_taps_to_literal"]

MODES = ("fp32", "bf16", "tf32")


def round_mode(x: torch.Tensor, mode: str) -> torch.Tensor:
    """``fp32``: unchanged; ``bf16``: round to nearest even bfloat16; ``tf32``: truncate to 10 mantissa bits (the
    probe P12 truncation model of the device's fp32 matmul operands; the device measured 1.4-1.75x its error)."""
    if mode == "fp32":
        return x
    if mode == "bf16":
        return x.to(torch.bfloat16).to(torch.float32)
    if mode == "tf32":
        return (x.contiguous().view(torch.int32) & ~0x1FFF).view(torch.float32)
    raise ValueError(f"unknown rounding mode {mode!r}; one of {MODES}")


@dataclass(frozen=True)
class Emulation:
    """Operand rounding per site (``fp32`` = exact float32). ``conv_act`` / ``conv_weight``: inputs and weights
    (+ bias) of the 53 convs; ``point_act`` / ``point_weight``: inputs and weights of the 17 point / voxel linears;
    ``gather``: rows produced by the embedding gathers (frustum2pixel, pixel2point: ``ttnn.embedding`` tables are
    bf16-only, so ``bf16`` is the stock-op device); ``upsample``: the decoder's interpolation operands. ``overrides``:
    per-layer modes by layer name (e.g. ``{"head.logits": ("tf32", "fp32")}``: (act, weight))."""

    conv_act: str = "fp32"
    conv_weight: str = "fp32"
    point_act: str = "fp32"
    point_weight: str = "fp32"
    gather: str = "fp32"
    upsample: str = "fp32"
    overrides: Dict[str, tuple] = field(default_factory=dict)

    @classmethod
    def parse(cls, spec: str) -> "Emulation":
        """``"conv_act=bf16,point_act=tf32,gather=bf16"`` (unset sites stay fp32)."""
        kw = {}
        for part in filter(None, (p.strip() for p in spec.split(","))):
            k, _, v = part.partition("=")
            if k not in cls.__dataclass_fields__ or k == "overrides" or v not in MODES:
                raise ValueError(f"bad emulation item {part!r}")
            kw[k] = v
        return cls(**kw)

    @property
    def exact(self) -> bool:
        return all(getattr(self, k) == "fp32" for k in ("conv_act", "conv_weight", "point_act", "point_weight",
                                                       "gather", "upsample")) and not self.overrides

    def label(self) -> str:
        items = [f"{k}={getattr(self, k)}" for k in ("conv_act", "conv_weight", "point_act", "point_weight",
                                                      "gather", "upsample") if getattr(self, k) != "fp32"]
        items += [f"{k}:{v}" for k, v in sorted(self.overrides.items())]
        return ",".join(items) or "fp32"


def segmax_logstep(x: torch.Tensor, plan: FramePlan, rounds: Optional[int] = None) -> torch.Tensor:
    """Per-voxel max of the sorted, padded rows ``x`` ``(N_cap, C)`` -> ``(M_cap, C)``: ``rounds`` (default: just
    enough for the frame's longest segment) log-step rounds over the host shift tables, then a gather of every
    segment's last row (PLAN.md 1.3). Exact (max is order-independent)."""
    rounds = segmax_rounds(plan.max_segment_length) if rounds is None else int(rounds)
    for idx in segmax_shift_tables(plan.seg_id, rounds):
        x = torch.maximum(x, x[_t(idx.astype(np.int64))])
    return x[_t(plan.seg_last.astype(np.int64))]


def _with_zero_row(table: torch.Tensor) -> torch.Tensor:
    return torch.cat([table, table.new_zeros(1, table.shape[1])], 0)


def _nhwc(x: torch.Tensor) -> torch.Tensor:
    _, c, h, w = x.shape
    return x.permute(0, 2, 3, 1).reshape(h * w, c)


def _nchw(rows: torch.Tensor, h: int, w: int) -> torch.Tensor:
    return rows.reshape(1, h, w, -1).permute(0, 3, 1, 2).contiguous()


class FRNetPlanTorch:
    """``FRNetPlanTorch(weights)(plan, taps)`` -> ``pred_probs`` of the sorted, padded rows ``(N_cap, 27)``
    (``plan.unsort`` gives the network-input order). Taps use the names of ``reference/model.py``; point taps are in
    sorted padded order, voxel taps have ``M_cap`` rows (:func:`plan_taps_to_literal` aligns them)."""

    def __init__(self, weights: FRNetWeights, *, emulation: Optional[Emulation] = None,
                 rounds: Optional[int] = None):
        self.weights = weights
        self.layers = TorchLayers(weights)
        self.emu = emulation or Emulation()
        self.rounds = rounds
        H, W = weights.frustum_hw
        self._interp = {}
        for st in [1] + [s.stride for s in weights.stages]:
            ah = interp_matrix(H // st, H, mode="linear", align_corners=True)
            aw = interp_matrix(W // st, W, mode="linear", align_corners=True)
            self._interp[st] = (_t(ah.astype(np.float32)), _t(aw.astype(np.float32)))
        # K-split weights of the concat-consuming layers
        self._split_lin = {}
        for lin, parts in ((weights.point_stem, (128, 256)), (weights.point_fuse[0], (128,) * 5),
                           *((st.point_fusion, (128, 128)) for st in weights.stages)):
            w_t, b_t, _ = self.layers.lin[lin.name]
            if w_t.shape[0] != sum(parts):
                raise ValueError(f"{lin.name}: {w_t.shape[0]} inputs, expected {sum(parts)}")
            self._split_lin[lin.name] = (list(torch.split(w_t, list(parts), 0)), b_t)
        self._split_conv = {}
        for conv, parts in ((weights.fusion_stem, (128, 128)), (weights.fuse_layer1, (128,) * 5),
                            *((st.pixel_fusion, (128, 128)) for st in weights.stages)):
            w_t, b_t, stride, pad, act = self.layers.conv[conv.name]
            if w_t.shape[1] != sum(parts):
                raise ValueError(f"{conv.name}: {w_t.shape[1]} inputs, expected {sum(parts)}")
            self._split_conv[conv.name] = (list(torch.split(w_t, list(parts), 1)), b_t, stride, pad, act)

    # ---- rounded layer applications ---------------------------------------------------------------------------
    def _modes(self, name: str, act: str, weight: str) -> tuple:
        return self.emu.overrides.get(name, (act, weight))

    def linear(self, parts: Sequence[torch.Tensor], layer: Linear) -> torch.Tensor:
        act_m, w_m = self._modes(layer.name, self.emu.point_act, self.emu.point_weight)
        if layer.name in self._split_lin:
            ws, b = self._split_lin[layer.name]
        else:
            w, b, _ = self.layers.lin[layer.name]
            ws = [w]
        if len(parts) != len(ws):
            raise ValueError(f"{layer.name}: {len(parts)} input parts for {len(ws)} weight parts")
        y = round_mode(b, w_m).expand(parts[0].shape[0], -1).clone()
        for x, w in zip(parts, ws):
            y = y + round_mode(x, act_m) @ round_mode(w, w_m)
        return F.relu(y) if layer.relu else y

    def conv(self, parts: Sequence[torch.Tensor], layer: Conv) -> torch.Tensor:
        act_m, w_m = self._modes(layer.name, self.emu.conv_act, self.emu.conv_weight)
        if layer.name in self._split_conv:
            ws, b, stride, pad, act = self._split_conv[layer.name]
        else:
            w, b, stride, pad, act = self.layers.conv[layer.name]
            ws = [w]
        if len(parts) != len(ws):
            raise ValueError(f"{layer.name}: {len(parts)} input parts for {len(ws)} weight parts")
        y = None
        for i, (x, w) in enumerate(zip(parts, ws)):
            part = F.conv2d(round_mode(x, act_m), round_mode(w, w_m), round_mode(b, w_m) if i == 0 else None,
                            stride=stride, padding=pad)
            y = part if y is None else y + part
        if act == "hardswish":
            y = F.hardswish(y)
        elif act == "sigmoid":
            y = torch.sigmoid(y)
        return y

    def gather(self, table: torch.Tensor, index: np.ndarray) -> torch.Tensor:
        return round_mode(table, self.emu.gather)[_t(index.astype(np.int64))]

    def upsample(self, x: torch.Tensor, stride: int) -> torch.Tensor:
        if stride == 1:
            return x
        ah, aw = self._interp[stride]
        m = self.emu.upsample
        x = round_mode(x, m)
        return torch.einsum("hi,bcij,wj->bchw", round_mode(ah, m), x, round_mode(aw, m))

    # ---- the device graph -------------------------------------------------------------------------------------
    def __call__(self, plan: FramePlan, taps: TapRegistry = NULL_TAPS) -> np.ndarray:
        return self.forward(plan, taps)

    @torch.no_grad()
    def forward(self, plan: FramePlan, taps: TapRegistry = NULL_TAPS) -> np.ndarray:
        Wt = self.weights
        H, W = plan.frustum_hw
        if (H, W) != Wt.frustum_hw:
            raise ValueError(f"plan grid {plan.frustum_hw} != weights grid {Wt.frustum_hw}")

        def tap(name: str, value: torch.Tensor) -> torch.Tensor:
            if taps.wants(name):
                taps.tap(name, value.numpy())
            return value

        def to_image(table: torch.Tensor, s: int) -> torch.Tensor:
            return _nchw(self.gather(_with_zero_row(table), plan.pix2vox[s]), H // s, W // s)

        def to_points(img: torch.Tensor, s: int) -> torch.Tensor:
            return self.gather(_nhwc(img), plan.point2pix[s])

        x = tap("encoder.input", _t(plan.features))
        pf: List[torch.Tensor] = []
        for i, lin in enumerate(Wt.ffe):
            x = self.linear([x], lin)
            pf.append(tap(f"encoder.pf{i}", x))
        vmax = tap("encoder.voxel_max", segmax_logstep(x, plan, self.rounds))
        vfeat = tap("encoder.voxel_feats", self.linear([vmax], Wt.compression))
        img = tap("backbone.frustum", to_image(vfeat, 1))
        for c in Wt.stem:
            img = self.conv([img], c)
        tap("backbone.stem", img)
        p = tap("backbone.point_stem", self.linear([to_points(img, 1), pf[-1]], Wt.point_stem))
        seg = tap("backbone.segmax0", segmax_logstep(p, plan, self.rounds))
        x = tap("backbone.fusion_stem", self.conv([to_image(seg, 1), img], Wt.fusion_stem))
        maps, pts_out = [x], [p]
        for st in Wt.stages:
            li, s = st.index, st.stride
            for bi, (c1, c2, ds) in enumerate(st.blocks):
                y = self.conv([self.conv([x], c1)], c2)
                idt = self.conv([x], ds) if ds is not None else x
                x = tap(f"backbone.layer{li}.block{bi}", F.hardswish(y + idt))
            tap(f"backbone.layer{li}", x)
            p = tap(f"backbone.point_fusion{li}", self.linear([to_points(x, s), p], st.point_fusion))
            seg = tap(f"backbone.segmax{li}", segmax_logstep(p, plan, self.rounds))
            fuse = tap(f"backbone.pixel_fusion{li}", self.conv([to_image(seg, s), x], st.pixel_fusion))
            att = tap(f"backbone.attention{li}", self.conv([self.conv([fuse], st.attention[0])], st.attention[1]))
            x = tap(f"backbone.stage{li}", fuse * att + x)
            maps.append(x)
            pts_out.append(p)
        strides = [1] + [st.stride for st in Wt.stages]
        ups = [tap(f"decoder.up{k}", self.upsample(mp, s)) for k, (mp, s) in enumerate(zip(maps, strides))]
        y = tap("decoder.fuse1", self.conv(ups, Wt.fuse_layer1))
        y = tap("decoder.fuse2", self.conv([y], Wt.fuse_layer2))
        q = tap("decoder.point_fuse1", self.linear(pts_out, Wt.point_fuse[0]))
        q = tap("decoder.point_fuse2", self.linear([q], Wt.point_fuse[1]))
        h = to_points(y, 1)
        for i, (lin, skip) in enumerate(zip(Wt.head_mlps, (q, pf[2], pf[1], pf[0]))):
            h = tap(f"head.mlp{i}", self.linear([h], lin) + skip)
        logits = tap("head.logits", self.linear([h], Wt.conv_seg))
        return tap("pred_probs", torch.softmax(logits, -1)).numpy()


def plan_taps_to_literal(name: str, value: np.ndarray, plan: FramePlan) -> np.ndarray:
    """A tap of :class:`FRNetPlanTorch` in the layout of ``reference/model.py``: point tensors (``N_cap`` rows)
    unsorted to the network-input order, voxel tensors (``M_cap`` rows) cut to ``M``, maps unchanged."""
    v = np.asarray(value)
    if v.ndim == 2 and v.shape[0] == plan.n_cap:
        return plan.unsort(v)
    if v.ndim == 2 and v.shape[0] == plan.m_cap:
        return v[: plan.m]
    return v