# SPDX-License-Identifier: Apache-2.0
"""C22: heatmap peaks of the query-based LiDAR detectors -- sigmoid and the 3x3 local maximum (PLAN.md section 1.2
row C22; probe P6, common/PROBES.md).

TransFusion, BEVFusion and PTv3 keep a heatmap cell only where it is the maximum of its 3x3 neighbourhood, with
model-specific border and class rules:

- ``"transfusion"`` (TF, S:transfusion:169): ``MaxPool 3x3 s1 p0`` padded back with a ZERO border, every class: the
  border cells are always suppressed;
- ``"bevfusion"`` (BF, S:bevfusion:258): classes ``pooled_classes`` (default 0-3) as TF; the other classes are not
  pooled (``local_max = heat``): every cell passes;
- ``"ptv3"`` (PT, S:ptv3:217): ``local_max`` starts as the heat itself and only the interior of the pooled classes is
  replaced by the pool: border cells and the other classes pass.

All three are one exact masked chain (P6, bit-exact on the device, TF 95 / BF 85 / PT 145 us):

    keep = m_pass + (heat == max_pool2d(heat, 3x3, stride 1, padding 1)) * m_eq        out = heat * keep

with two constant masks ``m_eq`` / ``m_pass`` (disjoint, values 0 / 1, laid out as the NHWC rows ``[H*W, C]`` of the
heatmap). The pool's padding never matters: a padded window only covers border cells, where ``m_eq`` is 0. Not the
literal ``p0`` pool + pad form (222 us for TF and two more programs). ``ttnn.max_pool2d`` is bf16 only (FLOAT32 is
rejected, ``pool_op.cpp:30``), so the heat is bf16: :func:`sigmoid_heat` takes fp32 logits through an fp32 sigmoid and
rounds once.

Device layout: the heatmap is a TILE ``[1, 1, H*W, C]`` tensor (what a conv head writes); the result has the same
layout (the input of C21's class-major top-k row and of the query-heat gather of C26).

numpy only at import; ttnn is imported inside the device functions.
"""
from __future__ import annotations

from typing import Any, Dict, Optional, Sequence, Tuple

import numpy as np

__all__ = ["RULES", "DEFAULT_POOLED_CLASSES", "local_max_masks", "local_max_numpy", "local_max_literal_numpy",
           "sigmoid_heat", "sigmoid_heat_numpy", "LocalMax"]

RULES = ("transfusion", "bevfusion", "ptv3")
DEFAULT_POOLED_CLASSES = 4        # BF / PT: classes 0-3 pooled (the last class(es) pass)


def _pooled(classes: int, rule: str, pooled_classes: Optional[Sequence[int]]) -> np.ndarray:
    if rule not in RULES:
        raise ValueError(f"rule {rule!r}: expected one of {RULES}")
    if pooled_classes is None:
        pooled_classes = range(classes) if rule == "transfusion" else range(min(DEFAULT_POOLED_CLASSES, classes))
    p = np.zeros(classes, np.float32)
    for c in pooled_classes:
        if not 0 <= int(c) < classes:
            raise ValueError(f"pooled class {c} outside [0, {classes})")
        p[int(c)] = 1.0
    if rule == "transfusion" and not p.all():
        raise ValueError("the transfusion rule pools every class")
    return p


def local_max_masks(height: int, width: int, classes: int, *, rule: str = "transfusion",
                    pooled_classes: Optional[Sequence[int]] = None, kernel: int = 3) -> Tuple[np.ndarray, np.ndarray]:
    """The constant masks ``(m_pass, m_eq)`` of the chain for ``rule`` (module docstring), float32 NHWC rows
    ``[height * width, classes]``. ``kernel`` (odd): the pool size; the border is ``kernel // 2`` cells wide."""
    if kernel < 1 or kernel % 2 == 0:
        raise ValueError(f"kernel must be odd and positive, got {kernel}")
    r = kernel // 2
    if height <= 2 * r or width <= 2 * r:
        raise ValueError(f"a {height}x{width} map has no interior for a {kernel}x{kernel} pool")
    pooled = _pooled(classes, rule, pooled_classes)
    interior = np.zeros((height, width, 1), np.float32)
    interior[r:height - r, r:width - r] = 1.0
    m_eq = interior * pooled[None, None, :]
    if rule == "transfusion":
        m_pass = np.zeros_like(m_eq)
    elif rule == "bevfusion":
        m_pass = np.broadcast_to(1.0 - pooled[None, None, :], m_eq.shape).astype(np.float32)
    else:
        m_pass = 1.0 - m_eq
    return (np.ascontiguousarray(m_pass.reshape(height * width, classes)),
            np.ascontiguousarray(m_eq.reshape(height * width, classes)))


def _pool_numpy(heat_nhwc: np.ndarray, height: int, width: int, kernel: int) -> np.ndarray:
    """3x3 (k x k) stride-1 max pool with ``-inf`` padding of NHWC rows ``[H*W, C]`` (exact on any float data)."""
    c = heat_nhwc.shape[-1]
    x = np.asarray(heat_nhwc, np.float64).reshape(height, width, c)
    r = kernel // 2
    p = np.full((height + 2 * r, width + 2 * r, c), -np.inf)
    p[r:r + height, r:r + width] = x
    out = np.full_like(x, -np.inf)
    for dy in range(kernel):
        for dx in range(kernel):
            np.maximum(out, p[dy:dy + height, dx:dx + width], out=out)
    return out.reshape(height * width, c)


def local_max_numpy(heat_nhwc: Any, height: int, width: int, m_pass: Any, m_eq: Any, kernel: int = 3) -> np.ndarray:
    """Host oracle of the masked chain on NHWC rows ``[H*W, C]`` (float32 out; exact on bf16-valued input)."""
    h = np.asarray(heat_nhwc, np.float32).reshape(height * width, -1)
    pooled = _pool_numpy(h, height, width, kernel)
    keep = np.asarray(m_pass, np.float32) + (h == pooled).astype(np.float32) * np.asarray(m_eq, np.float32)
    return (h * keep).astype(np.float32)


def local_max_literal_numpy(heat_nchw: Any, rule: str = "transfusion", *,
                            pooled_classes: Optional[Sequence[int]] = None, kernel: int = 3) -> np.ndarray:
    """The models' literal forms (``heat * (heat == local_max)``, NCHW ``(C, H, W)`` or ``(1, C, H, W)``), the oracle
    the masked chain is proven against: TF ``max_pool p0`` + zero border (all classes); BF pooled classes as TF and
    the others ``local_max = heat``; PT ``local_max = heat`` with the interior of the pooled classes pooled."""
    x = np.asarray(heat_nchw, np.float32)
    squeeze = x.ndim == 4
    if squeeze:
        x = x[0]
    c, h, w = x.shape
    pooled = _pooled(c, rule, pooled_classes).astype(bool)
    r = kernel // 2
    full = _pool_numpy(x.reshape(c, h * w).T, h, w, kernel).T.reshape(c, h, w)
    interior = np.zeros((h, w), bool)
    interior[r:h - r, r:w - r] = True
    if rule == "transfusion":
        local = np.where(interior[None], full, 0.0)
    elif rule == "bevfusion":
        local = np.where(pooled[:, None, None], np.where(interior[None], full, 0.0), x)
    else:
        local = np.where(pooled[:, None, None] & interior[None], full, x)
    out = (x * (x == local)).astype(np.float32)
    return out[None] if squeeze else out


def sigmoid_heat_numpy(logits: Any, dtype: str = "bfloat16") -> np.ndarray:
    """Host model of :func:`sigmoid_heat`: float64 sigmoid of the float32 logits, rounded once to ``dtype`` (bf16 RNE)."""
    s = (1.0 / (1.0 + np.exp(-np.asarray(logits, np.float64)))).astype(np.float32)
    if dtype in ("bfloat16", "bf16"):
        from ..tensors import round_to_bf16

        return round_to_bf16(s)
    return s


def sigmoid_heat(logits: Any, *, dtype: str = "bfloat16", memory_config: Any = None):
    """Device ``sigmoid(logits)`` computed in the logits' dtype (fp32 logits: an fp32 sigmoid) and cast once to
    ``dtype`` (bf16: what ``max_pool2d`` and C21's ``topk_large_indices`` take). Returns ``(heat, sig)``: the cast
    heat and the sigmoid before the cast (the same tensor when no cast is needed)."""
    import ttnn

    from ..tensors import ttnn_dtype

    kw = {} if memory_config is None else {"memory_config": memory_config}
    sig = ttnn.sigmoid(logits, **kw)
    want = ttnn_dtype(dtype)
    heat = sig if sig.dtype == want else ttnn.typecast(sig, want)
    return heat, sig


def _mc(spec: Any):
    import ttnn

    if spec is None or spec == "dram":
        return ttnn.DRAM_MEMORY_CONFIG
    if spec == "l1":
        return ttnn.L1_MEMORY_CONFIG
    return spec


class LocalMax:
    """The device chain for one heatmap geometry: ``heat`` (bf16 TILE ``[1, 1, H*W, C]``) -> ``heat * keep``.

    The masks are uploaded (bf16 TILE, exact 0 / 1) by the first call, outside any capture; with the TF rule
    (``m_pass`` = 0) the chain is ``max_pool2d``, ``eq``, ``multiply`` (keep = eq * m_eq), ``multiply``; otherwise
    ``addcmul(m_pass, eq, m_eq)`` replaces the first multiply. ``memory_config``: where the pool and the masks live
    (``"dram"`` default; ``"l1"`` is P6's 95 us form for maps that fit)."""

    def __init__(self, height: int, width: int, classes: int, *, rule: str = "transfusion",
                 pooled_classes: Optional[Sequence[int]] = None, kernel: int = 3, memory_config: Any = "dram",
                 name: str = "local_max"):
        self.height, self.width, self.classes = int(height), int(width), int(classes)
        self.rule = rule
        self.kernel = int(kernel)
        self.name = name
        self.m_pass, self.m_eq = local_max_masks(self.height, self.width, self.classes, rule=rule,
                                                 pooled_classes=pooled_classes, kernel=self.kernel)
        self.passes = bool(self.m_pass.any())
        self.memory_config = memory_config
        self._dev: Optional[Dict[str, Any]] = None

    @property
    def cells(self) -> int:
        return self.height * self.width

    def _upload(self, device: Any) -> Dict[str, Any]:
        import torch
        import ttnn

        def up(a: np.ndarray):
            t = torch.from_numpy(np.ascontiguousarray(a.reshape(1, 1, self.cells, self.classes)))
            return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device,
                                   memory_config=_mc(self.memory_config))

        self._dev = {"m_eq": up(self.m_eq), "m_pass": up(self.m_pass) if self.passes else None}
        return self._dev

    def __call__(self, heat: Any):
        import ttnn

        shape = tuple(int(s) for s in heat.shape)
        if shape != (1, 1, self.cells, self.classes):
            raise ValueError(f"{self.name}: heat must be [1, 1, {self.cells}, {self.classes}], got {shape}")
        if heat.dtype != ttnn.bfloat16:
            raise ValueError(f"{self.name}: max_pool2d takes bf16 only (got {heat.dtype}): use sigmoid_heat")
        dev = self._dev or self._upload(heat.device())
        r = self.kernel // 2
        pooled = ttnn.max_pool2d(input_tensor=heat, batch_size=1, input_h=self.height, input_w=self.width,
                                 channels=self.classes, kernel_size=[self.kernel, self.kernel], stride=[1, 1],
                                 padding=[r, r], dilation=[1, 1], ceil_mode=False,
                                 memory_config=_mc(self.memory_config), output_layout=ttnn.TILE_LAYOUT)
        if tuple(int(s) for s in pooled.shape) != shape:
            full = pooled
            pooled = ttnn.slice(full, [0, 0, 0, 0], list(shape))
            ttnn.deallocate(full, False)
        eq = ttnn.eq(heat, pooled)
        ttnn.deallocate(pooled, False)
        if self.passes:
            keep = ttnn.addcmul(dev["m_pass"], eq, dev["m_eq"])
        else:
            keep = ttnn.multiply(eq, dev["m_eq"])
        ttnn.deallocate(eq, False)
        out = ttnn.multiply(heat, keep)
        ttnn.deallocate(keep, False)
        return out

    def numpy(self, heat_nhwc: Any) -> np.ndarray:
        """Host oracle of this chain on NHWC rows (bf16-valued input gives the device result bit for bit)."""
        return local_max_numpy(heat_nhwc, self.height, self.width, self.m_pass, self.m_eq, self.kernel)

    def release(self) -> None:
        import ttnn

        if self._dev is not None:
            for t in self._dev.values():
                if t is not None and t.is_allocated():
                    ttnn.deallocate(t)
            self._dev = None

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "grid": [self.height, self.width], "classes": self.classes, "rule": self.rule,
                "kernel": self.kernel, "pooled_classes": [int(c) for c in np.flatnonzero(self.m_eq.max(axis=0))],
                "passes": self.passes, "memory_config": str(self.memory_config)}