# SPDX-License-Identifier: Apache-2.0
"""Host image pre-processing (``host.preprocess``) against the research reference's OpenCV pipeline, the
public-data converter's decoded-image digests and the research golden network input.

    TT_VISIBLE_DEVICES=none python -m pytest -q -p fake_ttnn_plugin code/tt_bevformer/tests/test_host_preprocess.py
"""
from __future__ import annotations

import hashlib

import numpy as np
import pytest

from tt_bevformer.host import preprocess as hp
from tt_bevformer.host.inputs import load_frame
from tt_bevformer.reference import config as C
from tt_bevformer.tests import paths

cv2 = pytest.importorskip("cv2")


def _sha(a: np.ndarray) -> str:
    return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()


@pytest.fixture(scope="module")
def frame30():
    paths.need_pandaset()
    return load_frame(paths.PANDASET_SAMPLES / "pandaset_019_f30.json")


@pytest.mark.parametrize("frame", paths.SAMPLE_FRAMES)
def test_shipped_images_decode_like_the_converter(frame):
    """OpenCV decode + the node's cv::resize to 1600x900 give exactly the BGR images the public-data converter
    hashed (so the stored CPU outputs are bit-for-bit inputs of this host path)."""
    paths.need_pandaset()
    fr = load_frame(paths.PANDASET_SAMPLES / f"pandaset_019_f{frame}.json")
    want = fr.manifest["image_bgr_sha256_1600x900"]
    for cam, im in zip(C.CAMERA_ORDER, fr.images_bgr):
        assert im.shape == (1080, 1920, 3) and im.dtype == np.uint8
        out = hp.node_transport(im)
        assert out.shape == (900, 1600, 3)
        assert _sha(out) == want[cam], cam


def test_pil_decode_equals_opencv(frame30):
    """The port decodes with PIL (``host.preprocess.decode_bgr``, ``ttaw.io``); the converter and the node use OpenCV
    (``cv2.imdecode`` / cv_bridge, BGR): the same pixels on the shipped JPEGs."""
    for cam, im in zip(C.CAMERA_ORDER, frame30.images_bgr):
        raw = (paths.PANDASET_SAMPLES / frame30.manifest["images"][cam]).read_bytes()
        assert np.array_equal(cv2.imdecode(np.frombuffer(raw, np.uint8), cv2.IMREAD_COLOR), im), cam


@pytest.mark.parametrize("hw", [(1080, 1920), (720, 1280), (432, 768), (1200, 1920)])
def test_node_transport(hw):
    """The node's uint8 cv::resize to 1600x900 (the vendored ttaw.image_linear emulation) equals OpenCV."""
    rng = np.random.default_rng(hw[0])
    im = rng.integers(0, 256, (900, 1600, 3), dtype=np.uint8)
    assert hp.node_transport(im) is im                                   # 1600x900 passes through
    src = rng.integers(0, 256, hw + (3,), dtype=np.uint8)
    assert np.array_equal(hp.node_transport(src), cv2.resize(src, (1600, 900), interpolation=cv2.INTER_LINEAR))
    with pytest.raises(ValueError):
        hp.node_transport(im.astype(np.float32))


def test_preprocess_matches_research_reference(frame30):
    R = paths.import_research()
    if R is None:
        pytest.skip("research/bevformer not in this workspace")
    imgs = [hp.node_transport(im) for im in frame30.images_bgr]
    got = hp.preprocess_images(frame30.images_bgr, workers=4)
    assert got.shape == (1, 6, 3, 736, 1280) and got.dtype == np.float32
    assert np.array_equal(got, R.preprocess_images_bgr(imgs))
    out = np.full((1, 6, 3, 736, 1280), 3.0, np.float32)
    assert np.array_equal(hp.preprocess_images(frame30.images_bgr, out=out), got)


def test_sample0_golden_input():
    """The research golden ``image`` of the nuScenes key-frame sample0 (internal validation data), bit for bit."""
    s0, golden = paths.vendor_sample0(), paths.research_file("golden", "sample0_golden.npz")
    if s0 is None or golden is None:
        pytest.skip("sample0 or its research golden are not in this workspace")
    imgs = [cv2.imread(str(s0 / "imgs" / f"{c}.jpg"), cv2.IMREAD_COLOR) for c in C.CAMERA_ORDER]
    with np.load(golden) as z:
        want = z["image"]
    assert np.array_equal(hp.preprocess_images(imgs, workers=4), want)


def test_preprocess_refuses_bad_inputs(frame30):
    with pytest.raises(ValueError):
        hp.preprocess_images(frame30.images_bgr[:5])