# SPDX-License-Identifier: Apache-2.0
"""Fused host input path: RGB uint8 image -> the window-order fp32 im2col rows of the patch-embed matmul, in one
threaded pass, written straight into the persistent host buffers that are uploaded to the device.

Reference path (``app._to_pixel_values`` + ``fused_tables.im2col_window_order``):

    x = u8 (CHW, fp32) / 255;  x = F.interpolate(x, (1008, 1008), bilinear, align_corners=False);  x = (x - 0.5) / 0.5
    im2col: [nW = 9, 576 tokens, 608 cols], cols (c, kh, kw), 588 .. 607 zero

This module computes the same values with the same fp32 arithmetic as torch's CPU bilinear kernel (source index
``fma(scale, dst + 0.5, -0.5)``, per axis ``fma(t0, w0, t1 * w1)``, width first; the tables of ``maskpost._axis_table``)
and writes each value to its im2col position directly (no fp32 CHW image, no interpolate output, no im2col copy).
The work is split by window row ``band`` (output pixel rows 336 b .. 336 b + 335 = the 3 windows of one window
row = rows 1728 b .. 1728 b + 1727 of the im2col matrix), so band b can be uploaded while band b + 1 is computed.

``self_check`` compares the result bit for bit with the reference path on several image sizes; the caller falls back
to the reference path when it fails (another torch build) or when numba is missing.
"""
from __future__ import annotations

import numpy as np

try:
    import numba
    from numba import njit, prange, types
    from numba.extending import intrinsic
except ImportError:  # optional dependency
    numba = None

from tt_sam3.server.maskpost import _axis_table

PATCH = 14
WIN = 24  # tokens per window side
OUT = 1008  # model input side
BAND_ROWS = PATCH * WIN  # 336 output pixel rows per window row
N_BANDS = OUT // BAND_ROWS  # 3
K = 3 * PATCH * PATCH  # 588
K_PAD = 608

_F32 = np.float32
# u8 / 255 in fp32 (correctly rounded division, == torch's float32 x.div_(255.0) for all 256 values)
_LUT = (np.arange(256, dtype=_F32) / _F32(255.0)).astype(_F32)

if numba is not None:

    @intrinsic
    def _fmaf(typingctx, a, b, c):  # noqa: ARG001 - numba intrinsic signature
        sig = types.float32(types.float32, types.float32, types.float32)

        def codegen(context, builder, signature, args):  # noqa: ARG001
            return builder.fma(args[0], args[1], args[2])

        return sig, codegen

    @njit(parallel=True, nogil=True, cache=True)
    def _band_kernel(src, pstride, lut, h0, h1, hl0, hl1, w0, w1, wl0, wl1, band, out):
        """src: (H, W * pstride) uint8 rows (RGB in the first 3 bytes of each pixel); out: (3, 576, 608) fp32 band."""
        half = np.float32(0.5)
        for r in prange(BAND_ROWS):
            y = band * BAND_ROWS + r
            a0 = h0[y]
            a1 = h1[y]
            la0 = hl0[y]
            la1 = hl1[y]
            ti = r // PATCH  # token row inside the window
            kh = r - ti * PATCH
            row0 = src[a0]
            row1 = src[a1]
            for x in range(OUT):
                b0 = w0[x] * pstride
                b1 = w1[x] * pstride
                lb0 = wl0[x]
                lb1 = wl1[x]
                tx = x // PATCH
                kw = x - tx * PATCH
                ww = tx // WIN
                tok = ti * WIN + (tx - ww * WIN)
                for c in range(3):
                    # width first (two source rows), then height: fma(t0, w0, t1 * w1) per axis
                    t0 = _fmaf(lut[row0[b0 + c]], lb0, lut[row0[b1 + c]] * lb1)
                    t1 = _fmaf(lut[row1[b0 + c]], lb0, lut[row1[b1 + c]] * lb1)
                    v = _fmaf(t0, la0, t1 * la1)
                    out[ww, tok, c * (PATCH * PATCH) + kh * PATCH + kw] = (v - half) / half

    @njit(parallel=True, nogil=True, cache=True)
    def _band_kernel_identity(src, pstride, lut, band, out):
        """1008 x 1008 input: no interpolation (``_to_pixel_values`` skips F.interpolate)."""
        half = np.float32(0.5)
        for r in prange(BAND_ROWS):
            y = band * BAND_ROWS + r
            ti = r // PATCH
            kh = r - ti * PATCH
            row = src[y]
            for x in range(OUT):
                tx = x // PATCH
                kw = x - tx * PATCH
                ww = tx // WIN
                tok = ti * WIN + (tx - ww * WIN)
                for c in range(3):
                    out[ww, tok, c * (PATCH * PATCH) + kh * PATCH + kw] = (lut[row[x * pstride + c]] - half) / half


_NUM_THREADS = [8]


def configure(num_threads: int) -> None:
    """Thread count of the band kernels (numba's setting is per calling thread; ``BandWriter.write`` applies it)."""
    if numba is not None:
        _NUM_THREADS[0] = max(1, min(int(num_threads), numba.config.NUMBA_NUM_THREADS))


def available() -> bool:
    return numba is not None


def _rows_view(rgb: np.ndarray):
    """(H, W, 3|4) uint8 (possibly the RGB view of an RGBA buffer) -> ((H, W * pstride) uint8 rows, pstride)."""
    if rgb.dtype != np.uint8 or rgb.ndim != 3 or rgb.shape[2] < 3:
        raise ValueError("expected (H, W, 3) uint8")
    h, w = rgb.shape[:2]
    s0, s1, s2 = rgb.strides
    if s2 == 1 and s1 in (3, 4) and s0 >= w * s1:
        # rows of the decoded buffer as they are (RGBA pixels, PNG filter-byte row padding): a strided 2-D view that
        # covers up to the last pixel's blue byte (never past the buffer)
        return np.lib.stride_tricks.as_strided(rgb, ((h, (w - 1) * s1 + 3)), (s0, 1), writeable=False), s1
    a = np.ascontiguousarray(rgb[..., :3])
    return a.reshape(h, w * 3), 3


class BandWriter:
    """Prepared per image: ``write(band, out)`` fills one (3, 576, 608) band of the window-order im2col matrix."""

    def __init__(self, rgb: np.ndarray):
        self.rows, self.pstride = _rows_view(rgb)
        h, w = rgb.shape[:2]
        self.identity = (h, w) == (OUT, OUT)
        if not self.identity:
            self.ht = _axis_table(h, OUT)
            self.wt = _axis_table(w, OUT)

    def write(self, band: int, out: np.ndarray) -> None:
        if out.shape != (N_BANDS, WIN * WIN, K_PAD) or out.dtype != _F32:
            raise ValueError("band buffer must be fp32 (3, 576, 608)")
        if numba.get_num_threads() != _NUM_THREADS[0]:
            numba.set_num_threads(_NUM_THREADS[0])
        if self.identity:
            _band_kernel_identity(self.rows, self.pstride, _LUT, band, out)
        else:
            h0, h1, hl0, hl1 = self.ht
            w0, w1, wl0, wl1 = self.wt
            _band_kernel(self.rows, self.pstride, _LUT, h0, h1, hl0, hl1, w0, w1, wl0, wl1, band, out)


def reference_bands(rgb: np.ndarray) -> np.ndarray:
    """The served reference path (``_to_pixel_values`` + window-order im2col) -> (9, 576, 608) fp32."""
    import torch

    from tt_sam3.imaging import DecodedImage, to_pixel_values as _to_pixel_values
    from tt_sam3.tt.fused_tables import im2col_window_order

    px = _to_pixel_values(DecodedImage(rgb))
    out = torch.zeros(N_BANDS * N_BANDS, WIN * WIN, K_PAD)
    return im2col_window_order(px, PATCH, WIN, out=out, k_pad=K_PAD, window_order=True).numpy()


def fused_bands(rgb: np.ndarray) -> np.ndarray:
    out = np.zeros((N_BANDS * N_BANDS, WIN * WIN, K_PAD), dtype=_F32)
    bw = BandWriter(rgb)
    for b in range(N_BANDS):
        bw.write(b, out[N_BANDS * b : N_BANDS * (b + 1)])
    return out


def self_check(seed: int = 0) -> bool:
    """Bit-for-bit comparison with the reference path (values incl. the zero pad columns) on random images of several
    sizes (down / up / identity / mixed scales, RGB and the RGB view of RGBA)."""
    if numba is None:
        return False
    rng = np.random.default_rng(seed)
    try:  # numba JIT / cache errors (e.g. no writable cache location in a container) -> the reference path
        for h, w, ch in ((1080, 1920, 4), (1008, 1008, 3), (720, 1280, 3), (1500, 1100, 4), (333, 517, 3), (1008, 2000, 3)):
            a = rng.integers(0, 256, size=(h, w, ch), dtype=np.uint8)
            rgb = a[..., :3] if ch == 4 else a
            if not np.array_equal(fused_bands(rgb).view(np.uint32), reference_bands(rgb).view(np.uint32)):
                return False
            # the PNG fast path's layout: RGBA rows W * 4 + 1 bytes apart (filter byte)
            rows = rng.integers(0, 256, size=(h, w * 4 + 1), dtype=np.uint8)
            rgb = rows[:, 1:].reshape(h, w, 4)[..., :3]
            if not np.array_equal(fused_bands(rgb).view(np.uint32), reference_bands(rgb).view(np.uint32)):
                return False
    except Exception:  # noqa: BLE001
        return False
    return True