# SPDX-License-Identifier: Apache-2.0
"""The CPU reference (``tt_sceneseg.reference``) against onnxruntime on the deployed ONNX: the oracle of every gate.

    PYTHONPATH=code tools/research-venv/bin/python -m pytest -q code/tt_sceneseg/tests/test_reference_cpu.py

Runs in the research venv (onnxruntime is not a dependency of the bundle; skipped where it is absent, e.g. in the
tt-metal environment and the image). ORT runs ``SceneSeg_FP32.onnx`` as shipped with every tap of ``config.TAPS``
exposed as an extra output; the torch reference runs on the same input (the node's exact pre-processing):

- before the data-dependent decision: every tap and the logits at PCC >= 0.9999 (PLAN.md 0.3 item 1; measured
  >= 0.9999999995);
- after it (the per-pixel argmax): set-based, the class maps and the published masks agree on >= 99.99 % of the
  pixels (fp32 vs fp32 differ only where two logits tie to within the accumulation-order noise).

Also: the reference built from the upstream state_dict (BN folded in fp64) against ORT, the public frames of
research/sceneseg/public_data and the research ORT goldens, and the SPEC's stored golden inputs.
"""
from __future__ import annotations

import numpy as np
import pytest

pytest.importorskip("onnxruntime")
pytest.importorskip("onnx")
torch = pytest.importorskip("torch")

from tt_sceneseg.host.postprocess import class_map_from_logits, postprocess  # noqa: E402
from tt_sceneseg.host.preprocess import prepare_bgr, prepare_image  # noqa: E402
from tt_sceneseg.reference import assets, config  # noqa: E402
from tt_sceneseg.reference.model import SceneSegReference  # noqa: E402
from tt_sceneseg.reference.ort import OrtReference  # noqa: E402
from tt_sceneseg.tests.research import PUBLIC, SAMPLE, SPEC_REF_RUNS, public_ids  # noqa: E402
from tt_sceneseg.ttaw.golden import TapRegistry  # noqa: E402
from tt_sceneseg.ttaw.metrics import error_stats, pcc  # noqa: E402

PCC_MIN = 0.9999
AGREE_MIN = 0.9999
ONNX = assets.onnx_path(required=False)
pytestmark = pytest.mark.skipif(ONNX is None, reason="SceneSeg_FP32.onnx not found (set SCENESEG_WEIGHTS_DIR)")


@pytest.fixture(scope="module", autouse=True)
def _threads():
    old = torch.get_num_threads()
    torch.set_num_threads(4)
    yield
    torch.set_num_threads(old)


@pytest.fixture(scope="module")
def ref():
    return SceneSegReference.from_onnx(ONNX)


@pytest.fixture(scope="module")
def ort_taps():
    return OrtReference(ONNX, taps=config.TAPS, threads=4)


@pytest.fixture(scope="module")
def ort_plain():
    return OrtReference(ONNX, threads=4)


def _agree(a_logits, b_logits, hw):
    a, b = postprocess(hw, logits=a_logits), postprocess(hw, logits=b_logits)
    return (a.class_map == b.class_map).mean(), (a.mask == b.mask).mean()


@pytest.mark.parametrize("order", config.CHANNEL_ORDERS)
def test_every_tap_matches_onnxruntime(order, ref, ort_taps):
    prep = prepare_image(str(SAMPLE), channel_order=order)
    x = prep.network_input()
    want = ort_taps.run(x)
    rec = TapRegistry()
    got = ref.run_numpy(x, taps=rec)
    report = {}
    for name in config.TAPS:
        assert rec[name].shape == want[name].shape, (name, rec[name].shape, want[name].shape)
        report[name] = error_stats(rec[name], want[name])
    worst = min(report.items(), key=lambda kv: kv[1]["pcc"])
    print(f"{order}: worst tap {worst[0]} pcc={worst[1]['pcc']:.12f}; logits max|d|={report['logits']['max_abs']:.2e}")
    bad = {k: v["pcc"] for k, v in report.items() if not v["pcc"] >= PCC_MIN}
    assert bad == {}, f"taps below PCC {PCC_MIN}: {bad}"
    assert pcc(got["logits"], want["logits"]) >= PCC_MIN
    cls_agree, mask_agree = _agree(got["logits"], want["logits"], prep.src_hw)        # after the decision: set-based
    assert cls_agree >= AGREE_MIN and mask_agree >= AGREE_MIN
    # the tapped session computes the same logits as the plain one
    np.testing.assert_array_equal(want["logits"], OrtReference(ONNX, threads=4).run(x)["logits"])


def test_state_dict_reference_matches_onnxruntime(ort_plain):
    """The reference built from the upstream state_dict (re-hosted safetensors or the .pth; BN folded in fp64)."""
    sd = assets.state_dict_path(required=False)
    if sd is None:
        pytest.skip("the .pth / safetensors is not available")
    prep = prepare_image(str(SAMPLE))
    x = prep.network_input()
    got = SceneSegReference.from_state_dict(sd).run_numpy(x)["logits"]
    want = ort_plain.run(x)["logits"]
    assert pcc(got, want) >= PCC_MIN
    cls_agree, mask_agree = _agree(got, want, prep.src_hw)
    assert cls_agree >= AGREE_MIN and mask_agree >= AGREE_MIN


@pytest.mark.parametrize("fid", ["ps019_f40_front_camera", "c10k_0009", "tut_highway_rain_2"])
def test_public_frames(fid, ref, ort_plain):
    """research/sceneseg/public_data frames, both orders: the reference vs ORT on the exact input, and ORT on the
    exact input vs the research ORT golden (whose input was one ulp off: float32 division)."""
    from PIL import Image

    if fid not in public_ids():
        pytest.skip(f"{fid} not converted")
    with Image.open(PUBLIC / "inputs" / fid / "source_bgr.png") as im:
        bgr = np.ascontiguousarray(np.asarray(im.convert("RGB"))[:, :, ::-1])
    gold = np.load(PUBLIC / "golden" / fid / "ort_fp32.npz")
    for order in config.CHANNEL_ORDERS:
        prep = prepare_bgr(bgr, channel_order=order)
        x = prep.network_input()
        want = ort_plain.run(x)["logits"]
        got = ref.run_numpy(x)["logits"]
        assert pcc(got, want) >= PCC_MIN, (fid, order)
        assert min(_agree(got, want, prep.src_hw)) >= AGREE_MIN, (fid, order)
        assert pcc(want, gold[f"logits_{order}"]) >= 0.9999999, (fid, order)
        assert (class_map_from_logits(want) == gold[f"cls_{order}"]).mean() >= AGREE_MIN, (fid, order)


def test_spec_golden_inputs(ref):
    """The SPEC's stored golden tensors (research/sceneseg/ref_runs/highway_normal_{1,2}: the deployed and the
    RGB-bilinear inputs with their ORT logits): the reference on the stored inputs reproduces the stored logits."""
    for name in ("highway_normal_1", "highway_normal_2"):
        for tag in ("cpp_bgr", "cpp_rgb"):
            p = SPEC_REF_RUNS / name / f"{tag}_input_f32.npy"
            if not p.is_file():
                pytest.skip("research/sceneseg/ref_runs not present")
            got = ref.run_numpy(np.load(p))["logits"]
            want = np.load(SPEC_REF_RUNS / name / f"{tag}_logits_ort_f32.npy")
            assert pcc(got, want) >= PCC_MIN, (name, tag)
            assert (class_map_from_logits(got) == class_map_from_logits(want)).mean() >= AGREE_MIN, (name, tag)