# SPDX-License-Identifier: Apache-2.0
"""C23: ``ttnn.grid_sample`` helpers (PLAN.md section 1.2 row C23; probe P9, common/PROBES.md; first built by the
BEVDet port for AlignBEV and FPN_LSS).

Rules from probe P9 (``common/probes/g4-attention-gridsample-msda.md``), enforced or encoded here:

- **fp32 grids** for every data-dependent grid (bf16 grids cost PCC 0.9988 at 53 px and 0.992 at 128 px);
- **input C % 32 == 0** (C = 80 is rejected; zero pad channels stay exactly 0): :func:`padded_channels`,
  :func:`pad_channels`;
- **per-batch grids**: grid N must equal input N (no broadcast over the batch);
- **ROW_MAJOR, interleaved (or height-sharded)** input and grid, ``padding_mode="zeros"`` only;
- **never a FLOAT32 input together with ``fp32_dest_acc_en``**: wrong, non-deterministic output (strict xfail in P9);
  :func:`grid_sample` refuses it;
- **exact integer coordinates**: the kernel recovers pixel coordinates as ``x = g * W/2 + (W - 1)/2`` in fp32
  (``grid_sample_reader_common.hpp``, align_corners=False), so the align_corners=False normalisation
  ``g = (2x + 1) / W - 1`` returns integer x exactly whenever ``(2x + 1) / W`` is exact in fp32 (``W`` a power of
  two): an identity warp is then a bit-exact copy (P9: the literal align_corners=True normalisation leaves 2.8-7.6 %
  of values off by up to 1/32). Zero padding acts in pixel space, so the normalisation choice does not change what
  is sampled: callers state pixel coordinates (:func:`affine_grid`, :func:`resize_grid`) and always call the op with
  ``align_corners=False``;
- the four bilinear weights are computed in fp32 and **truncated** (not rounded) to bf16 by the kernel: a 0.33-0.6 %
  rrmse floor for fractional coordinates (plan risk R3; K5 with fp32 weights removes it). :func:`emulate_grid_sample`
  reproduces that arithmetic on the host (the device test's oracle).

Host functions are numpy only (grids are built on the host, per frame or once, and uploaded as fp32 ROW_MAJOR
tensors); :func:`grid_sample` is the device call with the checks above. DCNv2 as one grid_sample with K = 9
(``batch_output_channels=True``, BEVFormer / METEOR) and the fused MSDA wrapper are added by their ports.
"""
from __future__ import annotations

from typing import Any, Optional, Sequence, Tuple

import numpy as np

__all__ = ["OUTSIDE", "padded_channels", "pad_channels", "pixel_to_grid", "kernel_pixels", "affine_pixel_coords",
           "affine_grid", "resize_pixel_coords", "resize_grid", "emulate_grid_sample", "grid_sample"]

F32 = np.float32
# A normalised coordinate whose four taps fall outside any input (pixel -1.5 W - 0.5): the sample is exactly 0.
OUTSIDE = F32(-4.0)


def padded_channels(c: int, multiple: int = 32) -> int:
    """The channel count ``grid_sample`` accepts for ``c`` real channels (the next multiple of 32)."""
    return -(-int(c) // int(multiple)) * int(multiple)


def pad_channels(x: np.ndarray, multiple: int = 32) -> np.ndarray:
    """Zero-pad the last (channel) dimension of an NHWC array to a multiple of ``multiple`` (no copy when aligned)."""
    x = np.asarray(x)
    c = x.shape[-1]
    cp = padded_channels(c, multiple)
    if cp == c:
        return x
    out = np.zeros(x.shape[:-1] + (cp,), x.dtype)
    out[..., :c] = x
    return out


def pixel_to_grid(px: Any, size: int) -> np.ndarray:
    """Pixel coordinate (0 = centre of the first pixel) -> normalised coordinate for ``align_corners=False``:
    ``(2 px + 1) / size - 1`` in float32 (exact for integer ``px`` when ``size`` is a power of two)."""
    p = np.asarray(px, F32)
    return (F32(2.0) * p + F32(1.0)) / F32(size) - F32(1.0)


def kernel_pixels(g: Any, size: int) -> np.ndarray:
    """The pixel coordinate the device kernel recovers from a normalised coordinate (align_corners=False):
    ``g * size/2 + (size - 1)/2`` in float32."""
    g = np.asarray(g, F32)
    return g * F32(size / 2.0) + F32((size - 1) / 2.0)


def affine_pixel_coords(transforms: np.ndarray, out_hw: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
    """Per batch n and output pixel (h, w): ``ix = a w + b h + c``, ``iy = d w + e h + f`` in float32 with one
    rounding per operation (``transforms`` [N, 6] = (a, b, c, d, e, f), the BEVDet AlignBEV / BEVFormer-style pixel
    affine) -> (ix, iy) [N, H, W]."""
    t = np.asarray(transforms, F32).reshape(-1, 6)
    h, w = out_hw
    hh, ww = np.meshgrid(np.arange(h, dtype=F32), np.arange(w, dtype=F32), indexing="ij")
    a, b, c, d, e, f = (t[:, k, None, None] for k in range(6))
    ix = (a * ww + b * hh) + c
    iy = (d * ww + e * hh) + f
    return ix.astype(F32), iy.astype(F32)


def affine_grid(transforms: np.ndarray, in_hw: Tuple[int, int], out_hw: Optional[Tuple[int, int]] = None, *,
                valid: Optional[Sequence[bool]] = None) -> np.ndarray:
    """fp32 grid [N, H, W, 2] (x, y) sampling an ``in_hw`` input at :func:`affine_pixel_coords` (align_corners=False
    normalisation: identity transforms are bit-exact copies for power-of-two input sizes). Batches with
    ``valid[n] == False`` sample :data:`OUTSIDE` (exact zeros, e.g. a history slot that holds no frame yet)."""
    out_hw = tuple(in_hw) if out_hw is None else tuple(out_hw)
    ix, iy = affine_pixel_coords(transforms, out_hw)
    grid = np.stack([pixel_to_grid(ix, in_hw[1]), pixel_to_grid(iy, in_hw[0])], axis=-1).astype(F32)
    if valid is not None:
        v = np.asarray(valid, bool).reshape(-1)
        if v.size != grid.shape[0]:
            raise ValueError(f"valid has {v.size} entries for {grid.shape[0]} batches")
        grid[~v] = OUTSIDE
    return np.ascontiguousarray(grid)


def resize_pixel_coords(in_size: int, out_size: int, align_corners: bool = True) -> np.ndarray:
    """Source pixel of every output pixel of a 1-D bilinear resize (float64), torch's ``F.interpolate`` rules:
    align_corners=True ``i (in - 1) / (out - 1)``; False ``(i + 0.5) in / out - 0.5`` (not clamped: the zero padding
    of grid_sample differs from torch's edge clamp there, so prefer C18 ``Resize2d`` for align_corners=False)."""
    i = np.arange(out_size, dtype=np.float64)
    if align_corners:
        return i * (in_size - 1) / max(out_size - 1, 1)
    return (i + 0.5) * in_size / out_size - 0.5


def resize_grid(in_hw: Tuple[int, int], out_hw: Tuple[int, int], *, align_corners: bool = True,
                batch: int = 1) -> np.ndarray:
    """Static fp32 grid [batch, Ho, Wo, 2] of a bilinear resize (source pixels in float64, normalised with the
    align_corners=False formula; call :func:`grid_sample`). For align_corners=True every source pixel lies inside
    the input, so the result equals ``F.interpolate(..., align_corners=True)`` up to the bf16 weight truncation."""
    sy = resize_pixel_coords(in_hw[0], out_hw[0], align_corners)
    sx = resize_pixel_coords(in_hw[1], out_hw[1], align_corners)
    gx = ((2.0 * sx + 1.0) / in_hw[1] - 1.0).astype(F32)
    gy = ((2.0 * sy + 1.0) / in_hw[0] - 1.0).astype(F32)
    grid = np.empty((out_hw[0], out_hw[1], 2), F32)
    grid[..., 0] = gx[None, :]
    grid[..., 1] = gy[:, None]
    return np.ascontiguousarray(np.broadcast_to(grid, (int(batch),) + grid.shape))


def _truncate_bf16(a: np.ndarray) -> np.ndarray:
    bits = np.ascontiguousarray(a, F32).view(np.uint32) & np.uint32(0xFFFF0000)
    return bits.view(F32)


def emulate_grid_sample(x: np.ndarray, grid: np.ndarray) -> np.ndarray:
    """Host emulation of the device's bilinear ``grid_sample`` (align_corners=False, zero padding): pixel recovery
    as :func:`kernel_pixels`, ``floor``, fp32 fractions and weights truncated to bf16, taps outside the input skipped.
    ``x`` [N, H, W, C], ``grid`` [N, Ho, Wo, 2] -> [N, Ho, Wo, C] float64 (the device sums the four bf16 products
    in its own order, so compare with a tolerance; an exact identity / outside sample is exact on both)."""
    x = np.asarray(x, np.float64)
    g = np.asarray(grid, F32)
    n, h, w, _ = x.shape
    px = kernel_pixels(g[..., 0], w)
    py = kernel_pixels(g[..., 1], h)
    x0 = np.floor(px).astype(np.int64)
    y0 = np.floor(py).astype(np.int64)
    fx = (px - x0.astype(F32)).astype(F32)
    fy = (py - y0.astype(F32)).astype(F32)
    one = F32(1.0)
    out = np.zeros(g.shape[:-1] + (x.shape[-1],), np.float64)
    bidx = np.arange(n)[:, None, None]
    for dy, wy in ((0, one - fy), (1, fy)):
        for dx, wx in ((0, one - fx), (1, fx)):
            yy, xx = y0 + dy, x0 + dx
            ok = (yy >= 0) & (yy < h) & (xx >= 0) & (xx < w)
            wgt = _truncate_bf16((wy * wx).astype(F32)).astype(np.float64) * ok
            vals = x[bidx, np.clip(yy, 0, h - 1), np.clip(xx, 0, w - 1)]
            out += vals * wgt[..., None]
    return out


def grid_sample(x: Any, grid: Any, *, compute_kernel_config: Any = None, memory_config: Any = None,
                batch_output_channels: bool = False):
    """``ttnn.grid_sample(x, grid)``: bilinear, zero padding, ``align_corners=False`` (pair it with
    :func:`pixel_to_grid` / :func:`affine_grid` / :func:`resize_grid`), with the P9 rules checked before the call.
    The output goes to DRAM by default."""
    import ttnn

    if x.layout != ttnn.ROW_MAJOR_LAYOUT or grid.layout != ttnn.ROW_MAJOR_LAYOUT:
        raise ValueError("grid_sample needs ROW_MAJOR input and grid")
    if int(x.padded_shape[-1]) % 32:
        raise ValueError(f"grid_sample input channels must be a multiple of 32, got {x.padded_shape[-1]} "
                         "(pad them: deform.pad_channels)")
    if int(x.shape[0]) != int(grid.shape[0]):
        raise ValueError(f"grid batch {grid.shape[0]} != input batch {x.shape[0]} (expand the grid per batch)")
    if grid.dtype != ttnn.float32:
        raise ValueError("use an fp32 grid for data-dependent sampling (probe P9: bf16 grids lose PCC)")
    if x.dtype == ttnn.float32 and compute_kernel_config is not None and \
            getattr(compute_kernel_config, "fp32_dest_acc_en", False):
        raise ValueError("a FLOAT32 input with fp32_dest_acc_en returns garbage (probe P9 defect)")
    return ttnn.grid_sample(x, grid, mode="bilinear", padding_mode="zeros", align_corners=False,
                            use_precomputed_grid=False, batch_output_channels=batch_output_channels,
                            memory_config=ttnn.DRAM_MEMORY_CONFIG if memory_config is None else memory_config,
                            compute_kernel_config=compute_kernel_config)