# SPDX-License-Identifier: Apache-2.0
"""Images -> the network input ``image [1, 6, 3, 736, 1280]`` float32 with the Autoware node's exact semantics
(research/bevformer/SPEC.md section 3.1):

1. ``cv_bridge`` BGR8; any image that is not 1600x900 goes through ``cv::resize(img, (1600, 900))`` (INTER_LINEAR on
   uint8, ``AW/src/bevformer_node.cpp:690-697``): :func:`node_transport`, the vendored bit-exact emulation
   ``ttaw.image_linear.opencv_resize_linear_u8`` (C14, ttaw 0.16.0; numpy, no OpenCV);
2. normalise FIRST in float (BGR mean, std 1), ``cv::resize`` x0.8 INTER_LINEAR on the float image, zero pad bottom /
   right to /32, CHW: the vendored C14 preset ``ttaw.image.bevformer_preprocess`` (bit-exact against OpenCV and the
   research golden input, ``common/tests/host/test_image_bevformer_host.py``).

The cameras must already be in the network order (``reference.config.CAMERA_ORDER``). numpy only; JPEG / PNG decoding
(:func:`decode_bgr`, and the API's ``ttaw.io``) uses PIL, which gives the same pixels as OpenCV's ``imdecode`` (what
``cv_bridge`` and the public-data converter use) on the shipped sample (``tests/test_host_preprocess.py``).
"""
from __future__ import annotations

from typing import Any, Sequence

import numpy as np

from ..reference import config as C
from ..ttaw.image import bevformer_preprocess
from ..ttaw.image_linear import opencv_resize_linear_u8

__all__ = ["node_transport", "decode_bgr", "rgb_to_bgr", "preprocess_images"]


def node_transport(image_bgr: np.ndarray) -> np.ndarray:
    """The node's ``cv::resize`` to 1600x900 (INTER_LINEAR on uint8) when the size differs; identity otherwise."""
    img = np.asarray(image_bgr)
    if img.ndim != 3 or img.shape[2] != 3 or img.dtype != np.uint8:
        raise ValueError(f"expected a uint8 (H, W, 3) BGR image, got {img.dtype} {img.shape}")
    if img.shape[:2] == C.NODE_HW:
        return img
    return opencv_resize_linear_u8(np.ascontiguousarray(img), C.NODE_HW)


def decode_bgr(data: bytes) -> np.ndarray:
    """JPEG / PNG bytes -> BGR uint8 (what ``cv_bridge`` ``bgr8`` hands the node), decoded with PIL."""
    from ..ttaw.io import decode_image

    return rgb_to_bgr(decode_image(bytes(data)))


def rgb_to_bgr(image_rgb: np.ndarray) -> np.ndarray:
    """``ttaw.io.load_image`` returns RGB; the network wants BGR (``to_rgb`` false)."""
    return np.ascontiguousarray(np.asarray(image_rgb)[:, :, ::-1])


def preprocess_images(images_bgr: Sequence[Any], *, workers: int = 1, out: np.ndarray = None) -> np.ndarray:
    """Six BGR uint8 images in network order -> ``image [1, 6, 3, 736, 1280]`` float32 (``out`` may be a
    preallocated array of that shape)."""
    imgs = [node_transport(im) for im in images_bgr]
    if len(imgs) != C.NUM_CAMS:
        raise ValueError(f"need {C.NUM_CAMS} images in the order {list(C.CAMERA_ORDER)}, got {len(imgs)}")
    target = None if out is None else out.reshape((C.NUM_CAMS, 3) + C.IMG_HW)
    res = bevformer_preprocess(imgs, scale=C.SCALE, pad_divisor=C.PAD_DIVISOR, out=target, workers=workers)
    return res.reshape((1, C.NUM_CAMS, 3) + C.IMG_HW)