# SPDX-License-Identifier: Apache-2.0
"""Host post-processing (``tt_sceneseg.host.postprocess`` / ``host.output``) against the deployed node's rule and the
research goldens (no device, no weights).

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sceneseg/tests/test_host_postprocess.py

Oracles: a literal transliteration of ``run_model_node.cpp:148-188`` (here) and of the research's independent one
(``research/sceneseg/scripts/verify_ort.py`` ``post_cpp``, when onnxruntime imports); ``numpy.argmax`` + ``cv2.resize``
(INTER_NEAREST); the SPEC's deployed-path goldens (``ref_runs/<image>/cpp_{bgr,rgb}_mask_srcres.png`` from the stored
ORT logits, bit for bit) and the public-data goldens (``golden/<id>/ort_fp32.npz`` -> ``cls_*`` and ``mask_*.png``,
bit for bit).
"""
from __future__ import annotations

import base64
import io
import json

import numpy as np
import pytest

from tt_sceneseg.host.output import SceneSegMask, to_output
from tt_sceneseg.host.postprocess import (class_map_from_logits, class_map_from_margins, foreground_mask,
                                          margins_from_logits, postprocess, to_source)
from tt_sceneseg.reference import config
from tt_sceneseg.tests.research import PUBLIC, SPEC_REF_RUNS, TUTORIAL, public_ids, research_module


def _scalar_node(logits: np.ndarray) -> np.ndarray:
    """run_model_node.cpp:148-163, one pixel at a time (small arrays only)."""
    c, h, w = logits[0].shape
    out = np.zeros((h, w), np.uint8)
    for y in range(h):
        for x in range(w):
            best, mx = 0, np.float32(-1e9)
            for k in range(c):
                s = np.float32(logits[0, k, y, x])
                if s > mx:
                    mx, best = s, k
            out[y, x] = best
    return out


def _logits(rng, shape=(1, 3, 32, 48), ties=True):
    lg = rng.normal(size=shape).astype(np.float32) * 4
    if ties:                                    # exact ties of every kind: the lowest index must win
        lg[0, 1, :4] = lg[0, 0, :4]
        lg[0, 2, 4:8] = lg[0, 1, 4:8]
        lg[0, 2, 8:12] = lg[0, 0, 8:12]
        lg[0, :, 12:14] = 1.5
        lg[0, 1, 14:16] = np.maximum(lg[0, 0, 14:16], lg[0, 2, 14:16])
    return lg


def test_class_decision_matches_the_node_and_argmax():
    rng = np.random.default_rng(0)
    for _ in range(3):
        lg = _logits(rng)
        want = _scalar_node(lg)
        got = class_map_from_logits(lg)
        np.testing.assert_array_equal(got, want)
        np.testing.assert_array_equal(got, np.argmax(lg[0], axis=0))       # first maximum == lowest index
        np.testing.assert_array_equal(class_map_from_margins(*margins_from_logits(lg)), want)
        m01, m12, _ = margins_from_logits(lg)
        np.testing.assert_array_equal(class_map_from_margins(m01, m12) == 1, want == 1)   # fg needs no m20
    assert got.dtype == np.uint8 and set(np.unique(got)) <= {0, 1, 2}
    # the -1e9 start of the scan: scores at or below it never win (argmax would differ; never seen in practice)
    lg = np.full((1, 3, 2, 2), -2e9, np.float32)
    lg[0, 2, 0, 0] = -1.5e9
    assert class_map_from_logits(lg)[0, 0] == 0 == _scalar_node(lg)[0, 0]


def test_research_post_cpp():
    """research verify_ort.post_cpp (the verification pass's independent C++ re-implementation)."""
    pytest.importorskip("onnxruntime")
    cv2 = pytest.importorskip("cv2")
    ver = research_module("verify_ort")
    rng = np.random.default_rng(1)
    for hw in ((414, 727), (1080, 1920), (100, 150)):
        lg = _logits(rng, (1, 3, 320, 640))
        best, mask = ver.post_cpp(lg, hw)
        res = postprocess(hw, logits=lg)
        np.testing.assert_array_equal(res.class_map, best)
        np.testing.assert_array_equal(res.mask, mask)
        np.testing.assert_array_equal(res.class_map_src, cv2.resize(best, (hw[1], hw[0]),
                                                                    interpolation=cv2.INTER_NEAREST))


def test_postprocess_equals_opencv_nearest():
    cv2 = pytest.importorskip("cv2")
    rng = np.random.default_rng(2)
    lg = _logits(rng, (1, 3, 320, 640))
    cls = np.argmax(lg[0], axis=0).astype(np.uint8)
    for hw in ((414, 727), (1080, 1920), (331, 681), (320, 640), (2160, 3840), (17, 29)):
        res = postprocess(hw, logits=lg)
        want = cv2.resize(np.where(cls == 1, 255, 0).astype(np.uint8), (hw[1], hw[0]), interpolation=cv2.INTER_NEAREST)
        np.testing.assert_array_equal(res.mask, want)
        assert res.mask.dtype == np.uint8 and res.mask.shape == hw and set(np.unique(res.mask)) <= {0, 255}
        np.testing.assert_array_equal(res.mask == 255, res.class_map_src == 1)
    np.testing.assert_array_equal(postprocess((414, 727), class_map=cls).mask, postprocess((414, 727), logits=lg).mask)
    with pytest.raises(ValueError):
        postprocess((10, 10))
    with pytest.raises(ValueError):
        postprocess((10, 10), class_map=cls[:100])
    with pytest.raises(ValueError):
        class_map_from_logits(np.zeros((2, 3, 4, 4), np.float32))
    np.testing.assert_array_equal(foreground_mask(np.array([[0, 1, 2]])), [[0, 255, 0]])
    np.testing.assert_array_equal(to_source(np.arange(6, dtype=np.uint8).reshape(2, 3), (2, 3)),
                                  np.arange(6).reshape(2, 3))


@pytest.mark.parametrize("name", ["highway_normal_1", "highway_normal_2"])
def test_spec_deployed_masks(name):
    """The SPEC's deployed-path goldens: the stored ORT logits of both orders -> the stored source-size masks."""
    from PIL import Image

    d = SPEC_REF_RUNS / name
    if not (d / "cpp_bgr_logits_ort_f32.npy").is_file():
        pytest.skip("research/sceneseg/ref_runs not present")
    for tag in ("cpp_bgr", "cpp_rgb"):
        lg = np.load(d / f"{tag}_logits_ort_f32.npy")
        with Image.open(d / f"{tag}_mask_srcres.png") as im:
            want = np.asarray(im)
        res = postprocess(want.shape, logits=lg)
        np.testing.assert_array_equal(res.mask, want, err_msg=f"{name} {tag}")
        with Image.open(d / f"{tag}_cls.png") as im:     # class map x 100, written by crosscheck.py from ITS ORT run:
            cls = np.asarray(im) // 100                     # an input one ulp away (float32 division), so one pixel
        assert (res.class_map == cls).mean() >= 0.99999     # of highway_normal_2 sits on a background / road tie


def test_public_goldens():
    """research/sceneseg/public_data/golden: ORT logits (both orders) of 23 frames -> the stored class maps and
    published masks, bit for bit."""
    from PIL import Image

    for fid in public_ids():
        d = PUBLIC / "golden" / fid
        if not (d / "ort_fp32.npz").is_file():
            pytest.skip("research/sceneseg/public_data/golden not present")
        with np.load(d / "ort_fp32.npz") as z:
            for order in config.CHANNEL_ORDERS:
                with Image.open(d / f"mask_{order}.png") as im:
                    want = np.asarray(im)
                res = postprocess(want.shape, logits=z[f"logits_{order}"])
                np.testing.assert_array_equal(res.class_map, z[f"cls_{order}"], err_msg=f"{fid} {order}")
                np.testing.assert_array_equal(res.mask, want, err_msg=f"{fid} {order}")


def test_output_body():
    rng = np.random.default_rng(4)
    res = postprocess((90, 160), logits=_logits(rng, (1, 3, 320, 640), ties=False))
    out = to_output(res, model="sceneseg-p150", channel_order="rgb")
    assert isinstance(out, SceneSegMask)
    body = json.loads(json.dumps(out.to_dict()))
    assert body["model"] == "sceneseg-p150" and body["frame_id"] == "camera"
    assert body["mask"]["shape"] == body["class_map"]["shape"] == [90, 160]
    assert body["class_names"] == list(config.CLASS_NAMES) and sum(body["class_counts"].values()) == 90 * 160
    assert body["meta"]["channel_order"] == "rgb" and body["meta"]["foreground_fraction"] == pytest.approx(
        float((res.mask == 255).mean()), abs=1e-6)
    from PIL import Image

    mask = np.asarray(Image.open(io.BytesIO(base64.b64decode(body["mask"]["data"]))))
    cmap = np.asarray(Image.open(io.BytesIO(base64.b64decode(body["class_map"]["data"]))))
    np.testing.assert_array_equal(mask, res.mask)
    np.testing.assert_array_equal(cmap, res.class_map_src)
    assert [d["label"] for d in out.to_dicts()] == [config.CLASS_NAMES[i] for i in np.unique(res.class_map_src)]
    no_map = to_output(res, model="m", channel_order="bgr", include_class_map=False).to_dict()
    assert "class_map" not in no_map and no_map["class_counts"] == body["class_counts"]
    npz = out.to_dict("npz")
    assert npz["mask"]["format"] == npz["class_map"]["format"] == "npz"
    with pytest.raises(ValueError):
        SceneSegMask(res.mask, class_map=res.class_map)             # 320x640 class map vs a 90x160 mask
    assert SceneSegMask(res.mask).class_counts().keys() <= {"0", "1"} | set(config.CLASS_NAMES)


def test_tutorial_class_maps_are_stable():
    """The stored small goldens of the 9 shipped samples (code/scripts/ref_golden.py): the host post-processing of
    the stored logits samples reproduces the stored class-map pixels."""
    from tt_sceneseg.tests.research import GOLDENS

    for name in TUTORIAL:
        path = GOLDENS / f"{name}.npz"
        if not path.is_file():
            pytest.skip("tests/goldens not generated (code/scripts/ref_golden.py)")
        with np.load(path) as z:
            for order in config.CHANNEL_ORDERS:
                yx = z["sample_yx"]
                lg = z[f"{order}.logits_samples"]                       # (3, K) at the sampled pixels
                cls = class_map_from_logits(lg[None, :, :, None])[:, 0]
                np.testing.assert_array_equal(cls, z[f"{order}.class_map"][yx[0], yx[1]])