# SPDX-License-Identifier: Apache-2.0
"""Host image pre-processing of the Autoware BEVDet node, exactly (S = research/bevdet/SPEC.md section 3).

The node path, per camera (``AW/src/bevdet_node.cpp:249-256``, ``AW/src/ros_utils.cpp:125-136``, vendor
``preprocess.cu`` + the ``bevdet::Preprocess`` plugin, S:184-209):

1. ``cv_bridge`` ``bgr8``; if the image is not 1600x900 it is ``cv::resize``d to 1600x900 (default
   ``INTER_LINEAR``). The node does NOT rescale the intrinsics (S:173); this port does by default
   (:func:`scale_intrinsics`, decision D8), and ``scale_intrinsics=False`` reproduces the quirk.
2. Planar **B, G, R** uint8 planes (``convert_RGBHWC_to_BGRCHW`` after a BGR->RGB swap): BEVDet was trained on BGR
   planes normalised with the "RGB" ImageNet mean / std (S:188).
3. ``bevdet::Preprocess``: nearest resize by r = 0.44f and crop (140, 0):
   ``src_row = roundf(i / r + 140 / r)``, ``src_col = roundf(j / r + 0 / r)`` (float32, round half away from zero),
   rows 318, 320, 323, ..., 898 and columns 0, 2, 5, 7, ..., 1598.
4. ``(x - mean[c]) / std[c]`` in float32 -> the network input [6, 3, 256, 704].

Split for the TT port: all four steps run on the host. Steps 1-3 produce the uint8 crop :func:`crop_u8` (3.2 MB per
frame) and step 4 (:func:`normalize`) the float32 network input, exactly as in the CPU reference; ``api.BEVDet``
uploads it as bf16 NHWC rows (6.5 MB per frame, PORT_LOG.md decision B14). Uploading the uint8 crop and normalising
on the device is optimisation work (OPT_BASELINE.md, opportunity 1). Steps 3-4 are the BEVDet preset of the shared
component C14 (``ttaw.image.bevdet_nearest_crop`` / ``bevdet_normalize``, added to ttaw 0.5.0 by this port); step 1's
resize needs OpenCV and stays here.
numpy (+ OpenCV only for the resize of non-1600x900 images); no torch, no ttnn.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Sequence, Tuple

import numpy as np

from ..reference.config import BEVDetConfig
from ..ttaw import image as timg

__all__ = ["gather_index", "node_transport", "crop_u8", "crop_from_rgb", "normalize", "scale_intrinsics",
           "PreparedImages", "prepare_images"]

F32 = np.float32


def _lut(cfg: BEVDetConfig) -> "timg.NearestCropLUT":
    lut = timg.bevdet_nearest_lut(cfg.src_hw, cfg.input_hw, cfg.crop_hw)
    if lut.resize != float(cfg.resize):
        raise ValueError(f"C14 ratio {lut.resize} != the configuration's {float(cfg.resize)}")
    return lut


def gather_index(cfg: BEVDetConfig) -> Tuple[np.ndarray, np.ndarray]:
    """``(rows [256], cols [704])`` the Preprocess plugin reads from the 1600x900 planes (float32 arithmetic):
    ``roundf(i / r + crop_h / r)``, ``roundf(j / r + crop_w / r)`` (C14 ``ttaw.image.bevdet_nearest_lut``)."""
    lut = _lut(cfg)
    return lut.rows, lut.cols


def _resize_linear(img: np.ndarray, hw: Tuple[int, int]) -> np.ndarray:
    """``cv::resize(img, (w, h))`` with the default INTER_LINEAR (what the node does for non-1600x900 images)."""
    try:
        import cv2
    except ImportError as e:  # pragma: no cover - cv2 is a declared dependency of the image
        raise RuntimeError("images that are not 1600x900 need OpenCV (opencv-python-headless) for the node's "
                           "cv::resize") from e
    h, w = hw
    return cv2.resize(np.ascontiguousarray(img), (w, h), interpolation=cv2.INTER_LINEAR)


def node_transport(image: np.ndarray, cfg: BEVDetConfig, *, channels: str = "rgb") -> np.ndarray:
    """One camera image (H, W, 3) uint8 -> the planar B, G, R uint8 planes [3, 900, 1600] the node hands to the
    network. ``channels``: ``"rgb"`` (``ttaw.io.load_image`` / ``decode_image`` output) or ``"bgr"`` (OpenCV)."""
    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.shape} {img.dtype}")
    if channels not in ("rgb", "bgr"):
        raise ValueError("channels must be 'rgb' or 'bgr'")
    bgr = img[:, :, ::-1] if channels == "rgb" else img
    if bgr.shape[:2] != cfg.src_hw:
        bgr = _resize_linear(bgr, cfg.src_hw)
    return np.ascontiguousarray(bgr.transpose(2, 0, 1))


def crop_u8(planes: np.ndarray, cfg: BEVDetConfig) -> np.ndarray:
    """[N, 3, 900, 1600] uint8 BGR planes -> [N, 3, 256, 704] uint8: the exact samples the network sees (the
    Preprocess gather before normalisation)."""
    planes = np.asarray(planes)
    if planes.ndim != 4 or planes.shape[1] != 3 or tuple(planes.shape[2:]) != cfg.src_hw:
        raise ValueError(f"expected [N, 3, {cfg.src_hw[0]}, {cfg.src_hw[1]}] planes, got {planes.shape}")
    return _lut(cfg).apply_planes(planes)


def crop_from_rgb(image: np.ndarray, cfg: BEVDetConfig, *, channels: str = "rgb") -> np.ndarray:
    """One (H, W, 3) image -> its [3, 256, 704] uint8 BGR crop without materialising the 1600x900 planes when no
    resize is needed (same result as ``crop_u8(node_transport(image)[None])[0]``)."""
    img = np.asarray(image)
    if img.ndim == 3 and img.shape[:2] == cfg.src_hw and img.dtype == np.uint8 and img.shape[2] == 3:
        return _lut(cfg).apply(img, channels=channels)
    return crop_u8(node_transport(img, cfg, channels=channels)[None], cfg)[0]


def normalize(crop: np.ndarray, cfg: BEVDetConfig) -> np.ndarray:
    """``(x - mean[c]) / std[c]`` in float32 (the vendor fp32 path) -> [N, 3, 256, 704] float32."""
    return timg.bevdet_normalize(crop, cfg.mean, cfg.std)


def scale_intrinsics(K: np.ndarray, sx: float, sy: float) -> np.ndarray:
    """Intrinsics of an image resized by (sx, sy) with ``cv::resize``'s pixel-centre alignment:
    ``fx' = sx fx``, ``cx' = sx (cx + 0.5) - 0.5`` (and the same for y). float64 3x3."""
    K = np.asarray(K, np.float64).copy()
    K[0, 0] *= sx
    K[1, 1] *= sy
    K[0, 1] *= sx
    K[0, 2] = sx * (K[0, 2] + 0.5) - 0.5
    K[1, 2] = sy * (K[1, 2] + 0.5) - 0.5
    return K


@dataclass
class PreparedImages:
    """The host half of the image path for one frame."""

    crop_u8: np.ndarray                  # [N, 3, 256, 704] uint8 B, G, R
    source_hw: List[Tuple[int, int]]     # (H, W) of each image as received
    scale: List[Tuple[float, float]]     # (sx, sy) of the node's resize to 1600x900 (1.0 when none)
    crop_hwc: Optional[np.ndarray] = None  # [N, 256, 704, 3] uint8 B, G, R rows (host.fast twin; crop_u8 is its view)


def prepare_images(images: Sequence[np.ndarray], cfg: BEVDetConfig, *, channels: str = "rgb",
                   out: Optional[np.ndarray] = None) -> PreparedImages:
    """N camera images (H, W, 3) uint8 in network camera order -> :class:`PreparedImages`. ``out`` may be a
    preallocated [N, 3, 256, 704] uint8 buffer (e.g. a ``ttaw.tensors.HostStaging`` view) written in place."""
    n = len(images)
    in_h, in_w = cfg.input_hw
    if out is None:
        out = np.empty((n, 3, in_h, in_w), np.uint8)
    elif out.shape != (n, 3, in_h, in_w) or out.dtype != np.uint8:
        raise ValueError(f"out must be [{n}, 3, {in_h}, {in_w}] uint8, got {out.shape} {out.dtype}")
    src_h, src_w = cfg.src_hw
    hw, scale = [], []
    for i, img in enumerate(images):
        img = np.asarray(img)
        hw.append((int(img.shape[0]), int(img.shape[1])))
        scale.append((src_w / img.shape[1], src_h / img.shape[0]))
        out[i] = crop_from_rgb(img, cfg, channels=channels)
    return PreparedImages(out, hw, scale)