# SPDX-License-Identifier: Apache-2.0
"""C19: row gathers with ``ttnn.embedding`` and the gather form of scatter (PLAN.md section 1.2 row C19; probe P4).

A scatter of ``P`` source rows into an ``N``-row destination (a BEV canvas, a to-dense map) is written as a
gather: the host builds ``index[n]`` = the source row of destination row ``n``, or ``sentinel`` (an explicit
all-zero row appended to the table) where nothing lands. Every shape is static, so it traces, and nothing has to
clear the destination first. Probe P4 (common/PROBES.md) measured it bit-exact at the real shapes, with:

- the table ROW_MAJOR bf16 ``[1, 1, V, D]`` (``D % 32 == 0`` for the fused tilized output), UINT32 indices
  ``[1, N]`` (ROW_MAJOR);
- **the default call** ``ttnn.embedding(index, table, layout=TILE, padding_idx=sentinel,
  embeddings_type=PADDED)``: 0.19-0.54 ms per canvas (CenterPoint 480 x 480: 0.21 ms); GENERIC is 5.4-14.6x slower
  on canvases (most indices hit the sentinel);
- **PADDED returns the table's own pad row, not zeros** (S:ptv3:304): the zero row must really be in the table
  (:func:`with_zero_rows`).

``ttnn.embedding`` is bf16-only; fp32 row gathers are ``ttnn.gather`` (bit-exact but ~1 us per index, probe P13),
not covered here. Indices must lie in ``[0, V)``: an out-of-range index reads outside the table (and can hang the
chip), so :func:`check_index` validates host index arrays before they are uploaded.

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

from typing import Any, Optional, Tuple

import numpy as np

__all__ = [
    "check_index",
    "gather_rows",
    "gather_rows_numpy",
    "scatter_rows",
    "scatter_rows_numpy",
    "with_zero_rows",
]


def check_index(index: Any, *, num_rows: int, sentinel: Optional[int] = None) -> np.ndarray:
    """Validate a host gather index and return it as a contiguous ``uint32`` array of shape ``[1, N]``.

    ``num_rows`` is the number of real source rows; valid entries are ``0 <= i < num_rows`` or ``i == sentinel``
    (the zero row, conventionally ``num_rows``). Raises ``ValueError`` for anything else (negative, too large,
    non-integral, non-finite), because an out-of-range index makes the device read outside the table."""
    a = np.asarray(index)
    if a.dtype.kind == "f":
        if not np.all(np.isfinite(a)) or np.any(a != np.floor(a)):
            raise ValueError("gather index holds non-integral or non-finite values")
    elif a.dtype.kind not in "iu":
        raise ValueError(f"gather index must be integral, got dtype {a.dtype}")
    flat = a.reshape(-1).astype(np.int64)
    valid = (flat >= 0) & (flat < int(num_rows))
    if sentinel is not None:
        valid |= flat == int(sentinel)
    bad = ~valid
    if np.any(bad):
        i = int(np.flatnonzero(bad)[0])
        raise ValueError(f"gather index [{i}] = {int(flat[i])} is outside [0, {num_rows})"
                         + ("" if sentinel is None else f" and is not the sentinel {sentinel}"))
    return np.ascontiguousarray(flat.astype(np.uint32).reshape(1, -1))


def gather_rows_numpy(index: Any, table: Any) -> np.ndarray:
    """Host oracle: ``out[n] = table[index[n]]`` (``table`` ``[V, D]``)."""
    t = np.asarray(table)
    return t.reshape(-1, t.shape[-1])[np.asarray(index).reshape(-1).astype(np.int64)]


def scatter_rows_numpy(rows: Any, index: Any, *, sentinel: int) -> np.ndarray:
    """Host oracle of :func:`scatter_rows`: ``rows`` ``[P, D]`` with a zero row appended at ``sentinel``
    (``sentinel >= P``; rows between ``P`` and ``sentinel`` are zero too), gathered by ``index``."""
    r = np.asarray(rows, dtype=np.float32)
    r = r.reshape(-1, r.shape[-1])
    table = np.zeros((max(int(sentinel) + 1, r.shape[0]), r.shape[1]), np.float32)
    table[: r.shape[0]] = r
    return gather_rows_numpy(index, table)


def _rm(t: Any):
    import ttnn

    return ttnn.to_layout(t, ttnn.ROW_MAJOR_LAYOUT) if t.layout != ttnn.ROW_MAJOR_LAYOUT else t


def with_zero_rows(rows: Any, count: int = 1):
    """A ROW_MAJOR gather table ``[1, 1, P + count, D]``: the device rows ``rows`` (``[1, 1, P, D]``, TILE or
    ROW_MAJOR, bf16) followed by ``count`` zero rows (the sentinel rows of :func:`scatter_rows`). One ``ttnn.pad``
    on the device (plus an untilize for a TILE input); traceable."""
    import ttnn

    if count < 1:
        raise ValueError("count must be >= 1")
    shape = tuple(int(s) for s in rows.shape)
    if len(shape) != 4 or shape[0] != 1 or shape[1] != 1:
        raise ValueError(f"rows must be [1, 1, P, D], got {shape}")
    rm = _rm(rows)
    return ttnn.pad(rm, [(0, 0), (0, 0), (0, int(count)), (0, 0)], 0.0)


def gather_rows(index: Any, table: Any, *, sentinel: Optional[int] = None, layout: str = "tile",
                memory_config: Any = None):
    """``out[0, 0, n] = table[0, 0, index[0, n]]`` on the device -> ``[1, 1, N, D]``.

    ``index``: UINT32 ROW_MAJOR ``[1, N]`` device tensor; ``table``: bf16 ROW_MAJOR ``[1, 1, V, D]`` (a TILE table
    is untilized by ttnn first: one extra program). ``sentinel``: the index of an all-zero table row; given, the
    call is P4's fast PADDED form (the table must really hold zeros there). ``layout="tile"`` (default) writes
    the tilized output directly when ``N % 32 == 0`` and ``D % 32 == 0``."""
    import ttnn

    lay = ttnn.TILE_LAYOUT if layout == "tile" else ttnn.ROW_MAJOR_LAYOUT
    kw = {"layout": lay}
    if sentinel is not None:
        kw.update(padding_idx=int(sentinel), embeddings_type=ttnn.EmbeddingsType.PADDED)
    if memory_config is not None:
        kw["memory_config"] = memory_config
    out = ttnn.embedding(index, table, **kw)               # [1, N, D]
    n, d = int(out.shape[-2]), int(out.shape[-1])
    return ttnn.reshape(out, (1, 1, n, d))


def scatter_rows(rows: Any, index: Any, *, sentinel: Optional[int] = None, layout: str = "tile",
                 memory_config: Any = None) -> Tuple[Any, Any]:
    """Gather-form scatter: ``canvas[n] = rows[index[n]]``, or zeros where ``index[n] == sentinel``.

    ``rows``: ``[1, 1, P, D]`` device rows (TILE or ROW_MAJOR bf16); ``sentinel`` defaults to ``P`` (the zero
    row appended by :func:`with_zero_rows`). Returns ``(canvas, table)``: the ``[1, 1, N, D]`` canvas (TILE by
    default, e.g. the NHWC input of a conv) and the table, which the caller may deallocate."""
    p = int(rows.shape[-2])
    sentinel = p if sentinel is None else int(sentinel)
    if sentinel < p:
        raise ValueError(f"sentinel {sentinel} must be >= the number of source rows {p}")
    table = with_zero_rows(rows, sentinel - p + 1)
    return gather_rows(index, table, sentinel=sentinel, layout=layout, memory_config=memory_config), table