# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing of the SceneSeg node, bit-exact (research/sceneseg/SPEC.md section 3; PLAN.md 2.2).

VisionPilot's ``scene_seg_model`` (``MR/ROS2/models/src/run_model_node.cpp:53-70``) converts the incoming image to
``bgr8`` with ``cv_bridge`` and its backend (``MR/common/backends/onnx_runtime_backend.cpp:41-60`` ==
``tensorrt_backend.cpp:160-177``) squashes it to 640x320 with ``cv::resize`` (INTER_LINEAR, aspect ratio not kept),
scales by 1/255, normalises with the ImageNet constants in B, G, R order and splits it into B, G, R planes. The vendored
``ttaw.image_linear`` (C14) reproduces every step bit for bit (OpenCV 4.8.1 and 4.11.0).

The TT port uploads the 640x320 uint8 B, G, R squash (614 KB; :attr:`PreparedImage.u8_bgr`) and folds the
normalisation and, for the ``rgb`` variant, the channel permutation into the device graph (PLAN 2.2: mean subtracted
before the stem's zero padding, 1 / (255 std) and the permutation folded into the stem weights). The CPU reference
consumes :meth:`PreparedImage.network_input`, the ONNX input exactly as the node builds it.

Channel order (LOAD-time variant): ``"bgr"`` is the deployed path (B, G, R planes into a network trained on R, G, B:
an upstream pre-processing bug the port keeps, user decision 1); ``"rgb"`` feeds R, G, B planes with the RGB
constants (training order). Each colour keeps its mean / std, so the two network inputs are channel permutations of
each other and both come from the same uint8 squash.

Two bit-identical implementations of the squash (knob ``SCENESEG_HOST_SQUASH``, OPT_PLAN.md H1): ``numba``
(default) is the compiled twin of ``host.fast`` with the RGB -> BGR swap folded into its reads (:func:`prepare_rgb`
takes the decoded RGB image directly); ``numpy`` is the vendored oracle on an explicit BGR copy.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional, Tuple

import numpy as np

from ..reference import config
from ..ttaw.image_linear import sceneseg_normalize, sceneseg_resize

__all__ = ["PreparedImage", "prepare_bgr", "prepare_rgb", "prepare_image", "rgb_to_bgr", "check_channel_order",
           "squash_method"]

SQUASH_METHODS = ("numba", "numpy")
_METHOD: Optional[str] = None


def squash_method() -> str:
    """The squash implementation of this process: knob ``SCENESEG_HOST_SQUASH`` (read once), ``numpy`` when numba is
    missing."""
    global _METHOD
    if _METHOD is None:
        from ..tt.knobs import KNOBS
        from . import fast

        m = KNOBS.read().HOST_SQUASH
        _METHOD = m if (m != "numba" or fast.HAVE_NUMBA) else "numpy"
    return _METHOD


def check_channel_order(channel_order: str) -> str:
    if channel_order not in config.CHANNEL_ORDERS:
        raise ValueError(f"channel_order must be one of {config.CHANNEL_ORDERS}, got {channel_order!r}")
    return channel_order


@dataclass(frozen=True)
class PreparedImage:
    """One camera frame after the node's pre-processing.

    ``src_hw``: the source image size (the published mask has it); ``u8_bgr``: the 640x320 squash, uint8
    ``(320, 640, 3)`` B, G, R (the device input, for both channel orders); ``channel_order``: the variant the
    network input is built for."""

    src_hw: Tuple[int, int]
    u8_bgr: np.ndarray
    channel_order: str = config.DEFAULT_CHANNEL_ORDER

    def network_input(self) -> np.ndarray:
        """The ONNX ``input`` float32 [1, 3, 320, 640] of this frame (``ttaw.image_linear.sceneseg_normalize``)."""
        return sceneseg_normalize(self.u8_bgr, channel_order=self.channel_order)

    def device_input(self) -> np.ndarray:
        """The uint8 canvas the TT trace uploads, ``[1, 1, 320 * 640, 3]`` (NHWC rows, B, G, R), contiguous."""
        h, w = config.INPUT_HW
        return np.ascontiguousarray(self.u8_bgr.reshape(1, 1, h * w, 3))


def rgb_to_bgr(rgb: Any) -> np.ndarray:
    """RGB uint8 HxWx3 (``ttaw.io`` decoding) -> B, G, R (``cv_bridge`` ``bgr8``), contiguous."""
    a = np.asarray(rgb)
    if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8:
        raise ValueError(f"expected an RGB uint8 HxWx3 image, got {a.dtype} {a.shape}")
    return np.ascontiguousarray(a[:, :, ::-1])


def prepare_bgr(bgr: Any, *, channel_order: str = config.DEFAULT_CHANNEL_ORDER,
                out: Optional[np.ndarray] = None, method: Optional[str] = None) -> PreparedImage:
    """A B, G, R uint8 HxWx3 frame of any size (the node's ``bgr8`` input) -> :class:`PreparedImage`.
    ``method``: ``numba`` / ``numpy`` (default: :func:`squash_method`), bit-identical."""
    check_channel_order(channel_order)
    a = np.asarray(bgr)
    if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8:
        raise ValueError(f"expected a BGR uint8 HxWx3 image, got {a.dtype} {a.shape}")
    if (method or squash_method()) == "numba":
        from . import fast

        sq = fast.squash_into(a, config.INPUT_HW, swap=False, out=out)
    else:
        sq = sceneseg_resize(a, out=out)
    return PreparedImage((int(a.shape[0]), int(a.shape[1])), sq, channel_order)


def prepare_rgb(rgb: Any, *, channel_order: str = config.DEFAULT_CHANNEL_ORDER,
                out: Optional[np.ndarray] = None, method: Optional[str] = None) -> PreparedImage:
    """A decoded R, G, B uint8 HxWx3 frame -> :class:`PreparedImage` = ``prepare_bgr(rgb_to_bgr(rgb))`` bit for bit.
    The ``numba`` squash reads the RGB image with the channel swap folded in (no BGR copy); ``numpy`` makes the
    copy and runs the oracle."""
    check_channel_order(channel_order)
    a = np.asarray(rgb)
    if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8:
        raise ValueError(f"expected an RGB uint8 HxWx3 image, got {a.dtype} {a.shape}")
    if (method or squash_method()) == "numba":
        from . import fast

        sq = fast.squash_into(a, config.INPUT_HW, swap=True, out=out)
        return PreparedImage((int(a.shape[0]), int(a.shape[1])), sq, channel_order)
    return prepare_bgr(rgb_to_bgr(a), channel_order=channel_order, out=out, method="numpy")


def prepare_image(image: Any, *, channel_order: str = config.DEFAULT_CHANNEL_ORDER) -> PreparedImage:
    """Any image the API takes (path, bytes, PIL image, RGB uint8 array, ``ttaw.io.CameraImage``) -> decoded to RGB by
    ``ttaw.io`` (alpha dropped, as ``cv::imdecode`` / ``cv_bridge`` do), then :func:`prepare_bgr`."""
    from ..ttaw import io as _io

    if isinstance(image, _io.CameraImage):
        rgb = image.image
    else:
        rgb = _io.load_image(image)
    return prepare_rgb(rgb, channel_order=channel_order)