# SPDX-License-Identifier: Apache-2.0
"""The CPU reference (``reference/``: ONNX initializers by node, BN folded where the port folds it, torch fp32) vs
ONNX Runtime on the deployed ONNX (PLAN.md 0.3 item 1; BUNDLE_CONVENTIONS.md section 10), on identical voxel tensors
(the host pre-processing of ``host/``):

- **before the data-dependent selection** every tap (decorated PFN input, PFN layer 0, PFN, canvas, every backbone
  conv, the deblocks, the neck concat, shared conv, heatmap head, the local-max heatmap, the input-independent key
  position embedding) at PCC >= 0.9999 against the same tensor of the deployed graph (exposed as an extra ORT output);
- **the top-500 selection set-based**: >= 99 % of the set shared, and every ORT proposal with heat >= 0.1 (the
  "strong proposals" of PLAN.md 2.5) in the reference's set (two fp32 implementations may swap low-score tail entries,
  S:transfusion:221);
- **after it teacher-forced**: the reference run with ORT's own top-k indices reproduces every downstream tap
  (query features / positions / heat, decoder layers, prediction heads, the three outputs) at PCC >= 0.9999,
  position by position;
- **end to end set-based**: the same detections after Autoware's post-processing (same label, centre within 1 cm,
  score within 1e-4) and the same counts at every stage;
- the plain ORT session (no extra outputs) reproduces the research's bitwise output goldens on their voxel tensors,
  and the reference reproduces the research's expected detections of the shipped PandaSet sample.

Cases: the shipped samples (PandaSet 019 f40, test.pcd), the Autoware sample-rosbag frame 150 (62,527 pillars: above
Autoware's 60,000, kept whole by N_MAX = 65,536; research data, license unstated) and the same frame through the
``autoware_compat`` shuffle emulation. Runs in the research venv (onnxruntime + torch, CPU only; threads <= 4)::

    cd bundles/transfusion-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python -m pytest -q -p no:cacheprovider \\
        code/tt_transfusion/tests/test_reference_cpu.py

Skipped where onnxruntime or the weights are absent.
"""
from __future__ import annotations

import json

import numpy as np
import pytest

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

from tt_transfusion.reference.model import SELECTION_TAPS  # noqa: E402
from tt_transfusion.reference.ort import ALL_TAPS  # noqa: E402
from tt_transfusion.reference.pipeline import TransFusionReference  # noqa: E402
from tt_transfusion.tests import _data  # noqa: E402
from tt_transfusion.ttaw.golden import TapRegistry  # noqa: E402
from tt_transfusion.ttaw.metrics import match_detections, pcc  # noqa: E402

PCC_MIN = 0.9999            # PLAN.md 0.3 item 1
THREADS = 4
STRONG = 0.1                # PLAN.md 2.5: overlap of strong proposals (heat >= 0.1) 100 %
CASES = {
    "pandaset_019_f40": ("sample", "pandaset_019_f40.npz", None),
    "test_pcd": ("sample", "test_pcd.npz", None),
    "rosbag_0150": ("research", "frame_0150", None),
    "rosbag_0150_compat": ("research", "frame_0150", {"mode": "autoware_compat", "offset": 12345}),
}


@pytest.fixture(scope="module")
def references():
    w = _data.weights_dir()
    torch.set_num_threads(THREADS)
    return (TransFusionReference(w, backend="torch", threads=THREADS),
            TransFusionReference(w, backend="ort", threads=THREADS, intermediates=True))


def _source(case):
    kind, name, sampling = CASES[case]
    src = _data.sample_path(name) if kind == "sample" else _data.research_cloud(name)
    return src, sampling


_RUNS = {}


def _runs(references, case):
    if case not in _RUNS:
        rt, ro = references
        src, sampling = _source(case)
        tt, to = TapRegistry(), TapRegistry()             # every tap (module and per layer)
        a = rt.run(src, sampling=sampling, taps=tt)
        b = ro.run(src, sampling=sampling, taps=to)
        _RUNS.clear()                                     # keep one case in memory (~1.5 GB of taps)
        _RUNS[case] = (a, b, tt, to)
    return _RUNS[case]


@pytest.mark.parametrize("case", list(CASES))
def test_taps_before_the_selection_match_onnxruntime(references, case):
    a, b, tt, to = _runs(references, case)
    assert np.array_equal(a.prepared.features, b.prepared.features)       # identical decorated features
    names = [n for n in ALL_TAPS if n not in SELECTION_TAPS]
    assert set(names) <= set(to.names()) and set(names) <= set(tt.names())
    worst = {}
    for name in names:
        x, y = tt[name], to[name]
        assert x.size == y.size, name
        worst[name] = pcc(x, y.reshape(x.shape))
    print(case, json.dumps({k: round(v, 10) for k, v in sorted(worst.items(), key=lambda kv: kv[1])[:5]}))
    bad = {k: v for k, v in worst.items() if not v >= PCC_MIN}
    assert not bad, bad


@pytest.mark.parametrize("case", list(CASES))
def test_topk_selection_set_based(references, case):
    a, b, tt, to = _runs(references, case)
    mine, ref = tt["head.topk_indices"].ravel(), to["head.topk_indices"].ravel()
    k = ref.size
    shared = len(set(mine.tolist()) & set(ref.tolist()))
    strong = ref[to["head.topk_values"].ravel() >= STRONG]
    print(case, f"top-{k} shared {shared}, strong {len(strong)}, positional {np.mean(mine == ref):.4f}")
    assert shared >= 0.99 * k
    assert set(strong.tolist()) <= set(mine.tolist())
    np.testing.assert_allclose(np.sort(tt["head.topk_values"].ravel()), np.sort(to["head.topk_values"].ravel()),
                               atol=1e-4)


@pytest.mark.parametrize("case", list(CASES))
def test_teacher_forced_decoder_matches_onnxruntime(references, case):
    rt, _ = references
    _, b, _, to = _runs(references, case)
    tf = TapRegistry(include=SELECTION_TAPS)
    rt.network(b.prepared, tf, topk_indices=to["head.topk_indices"])
    assert np.array_equal(tf["head.topk_indices"], to["head.topk_indices"])
    worst = {}
    for name in SELECTION_TAPS:
        if name == "head.topk_indices":
            continue
        x, y = tf[name], to[name]
        assert x.size == y.size, name
        worst[name] = pcc(x, y.reshape(x.shape))
    print(case, json.dumps({k: round(v, 10) for k, v in sorted(worst.items(), key=lambda kv: kv[1])[:5]}))
    bad = {k: v for k, v in worst.items() if not v >= PCC_MIN}
    assert not bad, bad
    assert np.array_equal(tf["query.pos"], to["query.pos"])               # exact arithmetic query_pos


@pytest.mark.parametrize("case", list(CASES))
def test_detections_match_onnxruntime_set_based(references, case):
    a, b, _, _ = _runs(references, case)
    assert a.counts == b.counts
    da, db = a.output, b.output
    m = match_detections(da.boxes, da.label_ids, da.scores, db.boxes, db.label_ids, db.scores, max_dist=0.01)
    assert m.recall == 1.0 and m.precision == 1.0, m.to_dict()
    assert max(m.score_errors, default=0.0) < 1e-4
    assert len(da) == len(db) and len(da) > 0


def test_compat_voxels_reproduce_the_research_shuffle_emulation(references):
    """``autoware_compat`` with offset 12345 = research ``run_ort_reference.py --autoware-shuffle`` (seed 0) on frame
    150: the voxel tensors of ``ort_sample_0150_awshuffle/io.npz`` bit for bit, and ORT's outputs on them."""
    a, b, _, _ = _runs(references, "rosbag_0150_compat")
    gold = np.load(_data.research_dir() / "ort_sample_0150_awshuffle" / "io.npz")
    mine = a.prepared.onnx_inputs()
    for k in ("voxels", "num_points", "coors"):
        assert np.array_equal(mine[k], gold[k]), k
    assert a.prepared.sampling["voxelizer_points"] == 2_000_000 and a.prepared.skipped is None


def test_plain_onnxruntime_reproduces_the_research_output_goldens(references):
    """The plain session (no extra outputs) on the research voxel tensors gives the research's bitwise ORT goldens
    (S:transfusion:221: public_data goldens and ``ort_sample_0250``)."""
    from tt_transfusion.reference.ort import OrtTransFusion

    rt, _ = references
    plain = OrtTransFusion(_data.weights_dir(), threads=THREADS, weights=rt.weights)
    r = _data.research_dir()
    cases = [(_data.pandaset_sample(), r / "public_data" / "goldens" / "ps019_f40" / "io.npz"),
             (_data.research_cloud("frame_0250"), r / "ort_sample_0250" / "io.npz")]
    for src, golden in cases:
        gold = np.load(golden)
        prep = rt.pre(src)
        mine = prep.onnx_inputs()
        for k in ("voxels", "num_points", "coors"):
            assert np.array_equal(mine[k], gold[k]), (golden, k)
        out = plain.run(mine)
        for k in ("cls_score0", "bbox_pred0", "dir_cls_pred0"):
            assert np.array_equal(out[k], gold[k]), (golden, k)


def test_pandaset_sample_reproduces_the_research_detections(references):
    """research/transfusion/public_data/shippable/pandaset_019_f40/expected_detections_cpu_fp32.json (ORT + the
    research post-processing, values rounded to 4-5 decimals): the same objects, labels after the remapper, scores and
    boxes."""
    from tt_transfusion.ttaw.decode import LABEL_IDS

    a, _, _, _ = _runs(references, "pandaset_019_f40")
    gold = json.loads((_data.research_dir() / "public_data" / "shippable" / "pandaset_019_f40" /
                       "expected_detections_cpu_fp32.json").read_text())["objects"]
    out = a.output
    assert len(out) == len(gold) == 80
    ref = np.array([[o[k] for k in ("x", "y", "z", "length", "width", "height", "yaw")] for o in gold], np.float64)
    labels = np.array([LABEL_IDS[o["label_remapped"]] for o in gold])
    m = match_detections(out.boxes, out.label_ids, out.scores, ref, labels, np.array([o["score"] for o in gold]),
                         max_dist=0.01)
    assert m.recall == 1.0 and m.precision == 1.0, m.to_dict()
    assert max(m.score_errors, default=0.0) < 1e-4
    for i, j in m.pairs:
        np.testing.assert_allclose(out.boxes[i, 2:6], ref[j, 2:6], atol=2e-4)