# SPDX-License-Identifier: Apache-2.0
"""C14 image pre-processing, StreamPETR preset: a bit-exact host emulation of Autoware's fused anti-aliased resize +
crop + normalise CUDA kernel (a PIL-style triangle filter with adaptive support), driven by cached per-size tables.
numpy only; no ttnn, no torch, no OpenCV.

Deployed code (autoware_universe main @ 9ceaccf, ``perception/autoware_camera_streampetr``; research/streampetr/SPEC.md
section 3.1):

- ``calculate_image_processing_params`` (``lib/network/camera_data_store.cpp:287-316``), float32:
  ``resize = max(480 / (float)H, 640 / (float)W)``; ``new_w = (int)(W * resize)``, ``new_h = (int)(H * resize)``;
  ``crop_h = max(0, (int)(1.0f * new_h) - 480)`` (the top rows are cropped, the bottom kept: ``bottom_crop_portion``
  0), ``crop_w = max(0, (new_w - 640) / 2)`` (centred) -> :func:`autoware_resize_geometry`.
- ``resize_and_extract_roi_kernel`` (``lib/network/preprocess_kernel.cu:73-176``; the 0.52.1 kernel
  ``resizeAndExtractRoi_kernel`` has the same arithmetic), per output pixel, all float32:
  ``r = (float)(out + roi_start)``, ``scale = (float)in / (float)resized``, ``centre = (r + 0.5f) * scale - 0.5f``,
  ``support = scale > 1 ? scale : 1``, taps ``ceilf(centre - support) .. floorf(centre + support)`` clamped to the
  image, ``w = max(0, 1 - |centre - i| / support)`` per axis, ``w = wx * wy``; the loop runs **y outer, x inner**
  and skips ``w <= 0``: ``sum_c += pix_c * w``, ``sum_w += w``; then ``sum_c /= sum_w`` (when ``sum_w > 0``) and
  ``out_c = (sum_c - mean[c]) / std[c]``, written planar (CHW). The channel order of the output is the byte order
  of the buffer the kernel reads: RGB on main (a ``bgr8`` message is swapped in place first,
  ``convert_bgr_to_rgb_kernel``), the message order on 0.52.1.

The emulation replays exactly these float32 operations in the same order, vectorised over the output pixels: per
output row / column the source taps and the 1-D weights are tabulated once (:class:`TriangleResizeLUT`), and the 2-D
accumulation runs tap pair by tap pair in the kernel's (y outer, x inner) order. Taps outside a pixel's window carry
weight 0, which adds exact zeros (the kernel skips them).

**Floating-point contraction** (``fma``). Written as C, every operation rounds once: ``fma=False`` (default) does
that, and equals the SPEC's numpy port (``research/streampetr/scripts/sp_common.py``) and the research goldens bit
for bit (their ``inputs_sha256.json``). The node's CMakeLists passes no ``--fmad`` / ``--use_fast_math`` flag, so
nvcc compiles with its default ``--fmad=true``, which lets the compiler fuse a multiply feeding an add into one
``fmaf``. ``fma=True`` models LLVM NVPTX's aggressive fusion: ``centre = fmaf(r + 0.5f, scale, -0.5f)``,
``sum_c = fmaf(pix_c, w, sum_c)`` and ``sum_w = fmaf(wx, wy, sum_w)`` (the rounded ``w`` still feeds the channel
products). Which form the deployed binary has cannot be checked without CUDA. The two differ in ~43 % of the
inputs by at most ~3e-5 in normalised units (one ulp of a ~1000 px centre moves every tap weight; measured in
``common/tests/host/test_image_triangle_host.py``), two orders of magnitude below the bf16 rounding of the device
input.

Normalisation presets (:data:`STREAMPETR_PRESETS`, SPEC section 3.4; PLAN.md D7 picks ``autoware_main``):

==================  ======================  ==================================================================
preset              planes (buffer order)   statistics
==================  ======================  ==================================================================
``autoware_main``   R, G, B                 mean (123.675, 116.28, 103.53), std (58.395, 57.12, 57.375)
                                            (universe main / 0.53.0, ``camera_data_store.cpp:128-134``)
``autoware_0.52``   R, G, B (an ``rgb8``    mean (103.53, 116.28, 123.675), std (57.375, 57.12, 58.395) on the
                    camera, as TIER IV's)   message bytes (0.52.1, ``camera_data_store.cpp:122-127``)
``awml_training``   B, G, R                 the 0.52 statistics (AWML t4base v2.5 test pipeline: cv2 decode,
                                            ``to_rgb=False``; = 0.52.1 fed a ``bgr8`` camera)
==================  ======================  ==================================================================
"""
from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Dict, Optional, Sequence, Tuple

import numpy as np

from .image import LUTCache

__all__ = [
    "TriangleGeometry",
    "TriangleResizeLUT",
    "StreamPETRPreset",
    "autoware_resize_geometry",
    "triangle_resize_lut",
    "fma_f32",
    "streampetr_preprocess",
    "streampetr_preprocess_batch",
    "STREAMPETR_PRESETS",
    "STREAMPETR_INPUT_HW",
    "STREAMPETR_MEAN_RGB",
    "STREAMPETR_STD_RGB",
    "STREAMPETR_MEAN_BGR",
    "STREAMPETR_STD_BGR",
]

F32 = np.float32
F64 = np.float64
STREAMPETR_INPUT_HW = (480, 640)   # HF ml_package_camera_streampetr.param.yaml:7-8 (input_image_height / _width)
# camera_data_store.cpp:128-134 (main: RGB order) and AW052 camera_data_store.cpp:122-127 (0.52.1: BGR order)
STREAMPETR_MEAN_RGB = (123.675, 116.280, 103.530)
STREAMPETR_STD_RGB = (58.395, 57.120, 57.375)
STREAMPETR_MEAN_BGR = (103.530, 116.280, 123.675)
STREAMPETR_STD_BGR = (57.375, 57.120, 58.395)


@dataclass(frozen=True)
class StreamPETRPreset:
    """One normalisation preset: the plane order the kernel sees (``"rgb"`` or ``"bgr"``) and the per-plane
    statistics, as float32 values (the kernel reads them from a float device buffer)."""

    name: str
    planes: str
    mean: Tuple[float, float, float]
    std: Tuple[float, float, float]
    description: str = ""

    def mean_f32(self) -> np.ndarray:
        return np.asarray(self.mean, dtype=F32)

    def std_f32(self) -> np.ndarray:
        return np.asarray(self.std, dtype=F32)


STREAMPETR_PRESETS: Dict[str, StreamPETRPreset] = {
    "autoware_main": StreamPETRPreset("autoware_main", "rgb", STREAMPETR_MEAN_RGB, STREAMPETR_STD_RGB,
                                      "Autoware universe main / 0.53.0: RGB planes, RGB-ordered statistics"),
    "autoware_0.52": StreamPETRPreset("autoware_0.52", "rgb", STREAMPETR_MEAN_BGR, STREAMPETR_STD_BGR,
                                      "Autoware 0.52.1 fed an rgb8 camera: message-order (RGB) planes, "
                                      "BGR-ordered statistics"),
    "awml_training": StreamPETRPreset("awml_training", "bgr", STREAMPETR_MEAN_BGR, STREAMPETR_STD_BGR,
                                      "AWML t4base v2.5 training / test pipeline (= 0.52.1 fed a bgr8 camera): "
                                      "BGR planes, BGR statistics"),
}


# --------------------------------------------------------------------------------------------- geometry


@dataclass(frozen=True)
class TriangleGeometry:
    """``calculate_image_processing_params`` of one source size: the virtual resized image and the ROI cut from it."""

    src_hw: Tuple[int, int]
    resized_hw: Tuple[int, int]       # (new_h, new_w)
    roi_hw: Tuple[int, int]           # the network input (480, 640)
    roi_start: Tuple[int, int]        # (start_y, start_x) in the resized image
    resize: float                     # the float32 resize factor, as an exact Python float

    @property
    def resize_f32(self) -> np.float32:
        return F32(self.resize)

    def to_dict(self) -> Dict[str, Any]:
        return {"src_hw": list(self.src_hw), "resized_hw": list(self.resized_hw), "roi_hw": list(self.roi_hw),
                "roi_start": list(self.roi_start), "resize": self.resize}


def autoware_resize_geometry(src_h: int, src_w: int, dst_h: int = STREAMPETR_INPUT_HW[0],
                             dst_w: int = STREAMPETR_INPUT_HW[1]) -> TriangleGeometry:
    """``camera_data_store.cpp:287-316`` in float32 (module docstring)."""
    src_h, src_w, dst_h, dst_w = int(src_h), int(src_w), int(dst_h), int(dst_w)
    if min(src_h, src_w, dst_h, dst_w) <= 0:
        raise ValueError(f"sizes must be positive: src {src_h}x{src_w}, dst {dst_h}x{dst_w}")
    scale_h = F32(F32(dst_h) / F32(src_h))
    scale_w = F32(F32(dst_w) / F32(src_w))
    resize = scale_w if scale_h < scale_w else scale_h            # std::max(scaleH, scaleW)
    new_w = int(F32(F32(src_w) * resize))                         # static_cast<int>(int * float): truncation
    new_h = int(F32(F32(src_h) * resize))
    crop_h = max(0, int(F32(F32(1.0) * F32(new_h))) - dst_h)      # (1.0f - bottom_crop_portion) * newH, portion 0
    crop_w = max(0, int((new_w - dst_w) / 2)) if new_w > dst_w else 0   # C int division, then std::max(0, .)
    return TriangleGeometry((src_h, src_w), (new_h, new_w), (dst_h, dst_w), (max(0, crop_h), max(0, crop_w)),
                            float(resize))


# --------------------------------------------------------------------------------------------- arithmetic


def fma_f32(a: Any, b: Any, c: Any) -> np.ndarray:
    """``fmaf(a, b, c)`` for float32 arrays: ``a * b + c`` rounded ONCE to float32 (round to nearest even), exactly.

    The float64 product of two float32 values is exact; the float64 sum may round, and a float64 result that lands
    on a float32 tie the exact value does not sit on would round twice. The TwoSum error term detects that case and
    the tie is broken towards the exact value."""
    a64 = np.asarray(a, dtype=F32).astype(F64)
    b64 = np.asarray(b, dtype=F32).astype(F64)
    c64 = np.asarray(c, dtype=F32).astype(F64)
    p = a64 * b64
    s = p + c64
    r = np.asarray(s.astype(F32))
    bv = s - p
    err = (p - (s - bv)) + (c64 - bv)            # p + c == s + err exactly (Knuth's TwoSum)
    bad = np.nonzero(np.broadcast_to(err, s.shape) != 0)
    if bad[0].size:
        sb, eb = s[bad], np.broadcast_to(err, s.shape)[bad]
        rb = r[bad]
        d = sb - rb.astype(F64)
        nb = np.nextafter(rb, np.where(d > 0, np.inf, -np.inf).astype(F32))
        tie = (d != 0) & (sb == (rb.astype(F64) + nb.astype(F64)) * 0.5)
        fix = tie & (np.sign(eb) == np.sign(d))
        rb = np.where(fix, nb, rb)
        r = r.copy()
        r[bad] = rb
    return r


def _axis_taps(n_out: int, start: int, n_in: int, n_resized: int, fma: bool) -> Tuple[np.ndarray, np.ndarray]:
    """Source index [n_out, T] and float32 weight [n_out, T] of every output row (or column), in the kernel's tap
    order; taps past a pixel's window have weight 0 (index clamped into the image)."""
    scale = F32(F32(n_in) / F32(n_resized))
    support = scale if scale > F32(1.0) else F32(1.0)
    r = (np.arange(n_out, dtype=np.int64) + int(start)).astype(F32)         # (float)(out + roi_start), exact
    r5 = (r + F32(0.5)).astype(F32)
    if fma:
        centre = (r5.astype(F64) * F64(scale) - 0.5).astype(F32)            # one rounding (exact in float64)
    else:
        centre = ((r5 * scale).astype(F32) - F32(0.5)).astype(F32)
    lo = np.maximum(0, np.ceil((centre - support).astype(F32)).astype(np.int64))
    hi = np.minimum(n_in - 1, np.floor((centre + support).astype(F32)).astype(np.int64))
    n_taps = max(1, int((hi - lo).max()) + 1)
    idx = lo[:, None] + np.arange(n_taps, dtype=np.int64)[None, :]
    valid = idx <= hi[:, None]
    idx = np.minimum(idx, n_in - 1)
    d = (centre[:, None] - idx.astype(F32)).astype(F32)
    w = (F32(1.0) - (np.abs(d) / support).astype(F32)).astype(F32)
    w = np.maximum(F32(0.0), w).astype(F32)
    w = np.where(valid, w, F32(0.0)).astype(F32)
    return idx, w


@dataclass(frozen=True)
class TriangleResizeLUT:
    """The per-size tables of ``resize_and_extract_roi_kernel``: build with :func:`triangle_resize_lut` (cached),
    apply with :meth:`apply` (the per-pixel weighted mean, before the mean / std normalisation)."""

    geometry: TriangleGeometry
    rows: np.ndarray       # (roi_h, Ty) int64 source rows
    row_w: np.ndarray      # (roi_h, Ty) float32 wy (0 past the window)
    cols: np.ndarray       # (roi_w, Tx) int64 source columns
    col_w: np.ndarray      # (roi_w, Tx) float32 wx
    fma: bool

    def apply(self, image: np.ndarray, *, out: Optional[np.ndarray] = None) -> np.ndarray:
        """One uint8 image ``(src_h, src_w, C)`` (any channel count, channel order kept) -> float32 ``(C, roi_h,
        roi_w)``: ``sum_c / sum_w`` (``sum_c`` where ``sum_w == 0``), exactly as the kernel before step 9."""
        g = self.geometry
        img = np.asarray(image)
        if img.ndim != 3 or img.dtype != np.uint8 or tuple(img.shape[:2]) != g.src_hw:
            raise ValueError(f"expected a uint8 ({g.src_hw[0]}, {g.src_hw[1]}, C) image, got {img.dtype} {img.shape}")
        roi_h, roi_w = g.roi_hw
        cn = img.shape[2]
        acc = np.zeros((cn, roi_h, roi_w), F32)
        wsum = np.zeros((roi_h, roi_w), F32)
        pix = np.empty((cn, roi_h, roi_w), F32)
        planes = np.ascontiguousarray(img.transpose(2, 0, 1))                   # (C, H, W) uint8
        # the source columns of every horizontal tap, once per image (float32 holds uint8 exactly)
        cols = [planes[:, :, self.cols[:, tx]].astype(F32) for tx in range(self.cols.shape[1])]
        for ty in range(self.rows.shape[1]):                                     # y outer ...
            rows = self.rows[:, ty]
            wy = self.row_w[:, ty][:, None]
            for tx in range(self.cols.shape[1]):                                 # ... x inner, as the kernel
                wx = self.col_w[:, tx][None, :]
                w = (wx * wy).astype(F32)                                        # w = wx * wy (rounded)
                np.take(cols[tx], rows, axis=1, out=pix)                         # (C, roi_h, roi_w) pixels
                if self.fma:
                    acc = fma_f32(pix, w, acc)                                   # sum_c = fmaf(pix, w, sum_c)
                    wsum = fma_f32(np.broadcast_to(wx, w.shape), np.broadcast_to(wy, w.shape), wsum)
                else:
                    np.multiply(pix, w, out=pix)                                 # pix * w (rounded)
                    np.add(acc, pix, out=acc)                                    # sum_c + . (rounded)
                    np.add(wsum, w, out=wsum)
        pos = wsum > 0
        result = np.where(pos, acc / np.where(pos, wsum, F32(1.0)), acc).astype(F32)
        if out is not None:
            if out.shape != result.shape or out.dtype != F32:
                raise ValueError(f"out must be float32 {result.shape}, got {out.dtype} {out.shape}")
            out[...] = result
            return out
        return result


_TRIANGLE_LUTS = LUTCache(maxsize=16)


def triangle_resize_lut(src_hw: Tuple[int, int], dst_hw: Tuple[int, int] = STREAMPETR_INPUT_HW, *,
                        geometry: Optional[TriangleGeometry] = None, fma: bool = False) -> TriangleResizeLUT:
    """The (cached) tables for one source size. ``geometry`` overrides the StreamPETR rule
    (:func:`autoware_resize_geometry`) for other nodes of the same kernel family (resized size, ROI start)."""
    g = geometry or autoware_resize_geometry(src_hw[0], src_hw[1], dst_hw[0], dst_hw[1])

    def make() -> TriangleResizeLUT:
        rows, row_w = _axis_taps(g.roi_hw[0], g.roi_start[0], g.src_hw[0], g.resized_hw[0], fma)
        cols, col_w = _axis_taps(g.roi_hw[1], g.roi_start[1], g.src_hw[1], g.resized_hw[1], fma)
        return TriangleResizeLUT(g, rows, row_w, cols, col_w, bool(fma))

    key = ("triangle", g.src_hw, g.resized_hw, g.roi_hw, g.roi_start, bool(fma))
    return _TRIANGLE_LUTS.get(key, make)


# --------------------------------------------------------------------------------------------- the preset


def _preset(preset: Any) -> StreamPETRPreset:
    if isinstance(preset, StreamPETRPreset):
        return preset
    try:
        return STREAMPETR_PRESETS[str(preset)]
    except KeyError:
        raise ValueError(f"unknown normalisation preset {preset!r}; one of {sorted(STREAMPETR_PRESETS)}") from None


def streampetr_preprocess(image: np.ndarray, *, preset: Any = "autoware_main", channels: str = "rgb",
                          dst_hw: Tuple[int, int] = STREAMPETR_INPUT_HW, fma: bool = False,
                          out: Optional[np.ndarray] = None) -> np.ndarray:
    """One camera image -> the float32 network input ``(3, 480, 640)`` of ``autoware_camera_streampetr``.

    ``image``: uint8 ``(H, W, 3)`` in the channel order ``channels`` (``"rgb"``: :func:`ttaw.io.load_image`, an
    ``rgb8`` message; ``"bgr"``: OpenCV, a ``bgr8`` message). The preset fixes the plane order the kernel sees and
    the statistics (:data:`STREAMPETR_PRESETS`); ``fma`` the contraction model (module docstring)."""
    p = _preset(preset)
    img = np.asarray(image)
    if img.ndim != 3 or img.shape[2] != 3 or img.dtype != np.uint8:
        raise ValueError(f"expected an (H, W, 3) uint8 image, got {img.dtype} {img.shape}")
    if channels not in ("rgb", "bgr"):
        raise ValueError("channels must be 'rgb' or 'bgr'")
    if channels != p.planes:
        img = img[:, :, ::-1]
    lut = triangle_resize_lut(img.shape[:2], dst_hw, fma=fma)
    v = lut.apply(img)                                                        # (3, H, W) float32
    mean = p.mean_f32()[:, None, None]
    std = p.std_f32()[:, None, None]
    res = ((v - mean).astype(F32) / std).astype(F32)                          # (sum_c - mean[c]) / std[c]
    if out is not None:
        if out.shape != res.shape or out.dtype != F32:
            raise ValueError(f"out must be float32 {res.shape}, got {out.dtype} {out.shape}")
        out[...] = res
        return out
    return res


def streampetr_preprocess_batch(images: Sequence[np.ndarray], *, preset: Any = "autoware_main",
                                channels: str = "rgb", dst_hw: Tuple[int, int] = STREAMPETR_INPUT_HW,
                                fma: bool = False, workers: int = 1, out: Optional[np.ndarray] = None) -> np.ndarray:
    """The cameras of one frame -> float32 ``(N, 3, 480, 640)`` (each camera may have its own size). ``workers``
    threads over the cameras (numpy releases the GIL in the array passes)."""
    imgs = list(images)
    shape = (len(imgs), 3) + tuple(int(v) for v in dst_hw)
    if out is None:
        out = np.empty(shape, F32)
    elif out.shape != shape or out.dtype != F32:
        raise ValueError(f"out must be float32 {shape}, got {out.dtype} {out.shape}")

    def one(i: int) -> None:
        streampetr_preprocess(imgs[i], preset=preset, channels=channels, dst_hw=dst_hw, fma=fma, out=out[i])

    if workers > 1 and len(imgs) > 1:
        with ThreadPoolExecutor(max_workers=min(int(workers), len(imgs))) as ex:
            list(ex.map(one, range(len(imgs))))
    else:
        for i in range(len(imgs)):
            one(i)
    return out