# SPDX-License-Identifier: Apache-2.0
"""The CPU reference (``reference/``: ONNX initializers by node, every BN folded where the port folds it, torch fp32)
vs the deployed graph's oracles (PLAN.md 0.3 item 1; BUNDLE_CONVENTIONS.md section 10), on identical inputs (the
host pre-processing of ``host/``). ONNX Runtime cannot load ``bevfusion_lidar.onnx`` (its sparse convs are TensorRT
plugins, S:bevfusion:343-351), so the oracle is split as the SPEC prescribes (SPEC 6.2, PLAN.md 2.6 "CPU reference"):

- **sparse encoder** vs the research reference (``research/bevfusion/scripts/bevfusion_ref.py``: its own hashed
  rulebooks and an ``index_add_`` scatter form with the BatchNorms unfolded, independent of this bundle's C13 gather
  maps): every layer output (conv_input, both blocks of every level, the three strided convs, conv_out) and the
  to-dense BEV at PCC >= 0.9999, the active sets of every level identical;
- **dense part before the data-dependent selection** vs ONNX Runtime on the deployed graph's plugin-free dense
  sub-graph fed with the same BEV: every tap (SECOND convs and blocks, deblocks, neck, shared conv, heatmap head,
  sigmoid heat, the local-max heatmap, the key position embedding) at PCC >= 0.9999;
- **the top-500 selection set-based**: >= 99 % of the set shared, and every ORT proposal with heat >= 0.1 (the
  strong proposals) in the reference's set;
- **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, ``bbox_pred``, ``score``) at PCC >= 0.9999
  and ``label_pred`` exactly;
- **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 research goldens (the research reference end to end: numpy voxelizer, torch scatter encoder, ORT dense): the
  same voxel tensors and BEV, the same heatmap, top-500 set and detections.

Cases: the shipped samples (PandaSet 019 f40, test.pcd), the Autoware sample-rosbag top scan #144 (the SPEC's frame;
research data, license unstated) and nuScenes mini_val scene-0103 key-frame 20 (sparse 32-beam; research only).
Runs in the research venv (onnxruntime + torch, CPU only; threads <= 4)::

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

Skipped where onnxruntime, the weights or the research reference 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_bevfusion.reference.model import SELECTION_TAPS  # noqa: E402
from tt_bevfusion.reference.ort import DENSE_ORT_TAPS  # noqa: E402
from tt_bevfusion.reference.pipeline import BEVFusionReference  # noqa: E402
from tt_bevfusion.reference.weights import ENCODER_PLAN  # noqa: E402
from tt_bevfusion.tests import _data  # noqa: E402
from tt_bevfusion.ttaw.golden import TapRegistry  # noqa: E402
from tt_bevfusion.ttaw.metrics import match_detections, pcc  # noqa: E402

PCC_MIN = 0.9999            # PLAN.md 0.3 item 1
THREADS = 4
STRONG = 0.1                # strong proposals (heat >= 0.1) must all be selected
CASES = {
    "pandaset_019_f40": ("sample", "pandaset_019_f40.npz"),
    "test_pcd": ("sample", "test_pcd.npz"),
    "rosbag_0144": ("research", "rosbag_0144"),
    "ns0103_k20": ("research", "ns0103_k20"),
}


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


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


_RUNS = {}


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


def _research_encoder_taps(R, inits, vfe, coors):
    """The research reference's sparse encoder (``SparseEncoderRef.forward``) layer by layer, with its own
    rulebooks, scatter-form convs and unfolded BatchNorms: {tap name: (N, C) features} + the level coords."""
    enc = R.SparseEncoderRef(inits)
    P = enc.P
    out, coords = {}, {}
    x = torch.from_numpy(vfe.astype(np.float32))
    c = torch.from_numpy(coors[:, ::-1].astype(np.int64).copy())
    level = R.SparseLevel(c, enc.SPATIAL)
    coords["l1"] = level.coords.numpy()
    rb = R.subm_rulebook(level, (3, 3, 3), (1, 1, 1))
    x = torch.relu(enc.bn(R.apply_rulebook(x, rb, enc.w[P + "conv_input.0.weight"], len(level.coords)),
                          P + "conv_input.1"))
    out["sparse.conv_input"] = x
    for stage, nblocks, down in enc.STAGES:
        rb = R.subm_rulebook(level, (3, 3, 3), (1, 1, 1))
        for blk in range(nblocks):
            pre = f"{P}encoder_layers.encoder_layer{stage}.{blk}."
            idn = x
            y = torch.relu(enc.bn(R.apply_rulebook(x, rb, enc.w[pre + "conv1.weight"], len(level.coords)),
                                  pre + "bn1"))
            out[f"sparse.layer{stage}.{blk}.conv1"] = y
            y = enc.bn(R.apply_rulebook(y, rb, enc.w[pre + "conv2.weight"], len(level.coords)), pre + "bn2")
            x = torch.relu(y + idn)
            out[f"sparse.layer{stage}.{blk}"] = x
        if down is not None:
            _, _, pad = down
            pre = f"{P}encoder_layers.encoder_layer{stage}.2."
            rb_d, level = R.strided_rulebook(level, (3, 3, 3), (2, 2, 2), pad)
            x = torch.relu(enc.bn(R.apply_rulebook(x, rb_d, enc.w[pre + "0.weight"], len(level.coords)), pre + "1"))
            out[f"sparse.down{stage}"] = x
            coords[f"l{stage + 1}"] = level.coords.numpy()
    rb_o, out_level = R.strided_rulebook(level, (1, 1, 3), (1, 1, 2), (0, 0, 0))
    x = torch.relu(enc.bn(R.apply_rulebook(x, rb_o, enc.w[P + "conv_out.0.weight"], len(out_level.coords)),
                          P + "conv_out.1"))
    out["sparse.conv_out"] = x
    coords["out"] = out_level.coords.numpy()
    return {k: v.numpy() for k, v in out.items()}, coords


@pytest.mark.parametrize("case", list(CASES))
def test_sparse_encoder_matches_the_research_reference(references, case):
    rt, _ = references
    a, _, tt, _ = _runs(references, case)
    R = _data.research_module()
    inits = R.load_initializers(str(_data.weights_dir() / _data.ONNX))
    ref, coords = _research_encoder_taps(R, inits, a.prepared.vfe, a.prepared.coors)
    plan = a.prepared.plan
    for lvl, c in coords.items():                       # identical active sets (the research's in key order)
        mine = plan.levels[lvl].coords
        assert len(mine) == len(c), lvl
        if lvl != "l1":
            assert np.array_equal(mine, c), lvl         # both sort strided outputs by the (x, y, z) linear key
    names = [name for name, *_ in ENCODER_PLAN]
    worst = {name: pcc(tt[name], ref[name]) for name in names}
    worst["sparse.bev"] = pcc(tt["sparse.bev"], R.SparseEncoderRef(inits).forward(a.prepared.vfe, a.prepared.coors))
    print(case, plan.describe()["counts"], json.dumps({k: round(v, 10) for k, v in
                                                       sorted(worst.items(), key=lambda kv: kv[1])[:4]}))
    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_dense_taps_before_the_selection_match_onnxruntime(references, case):
    a, b, tt, th = _runs(references, case)
    assert np.array_equal(a.prepared.vfe, b.prepared.vfe)
    names = [n for n in DENSE_ORT_TAPS if n not in SELECTION_TAPS]
    assert set(names) <= set(th.names()) and set(names) <= set(tt.names())
    worst = {}
    for name in names:
        x, y = tt[name], th[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):
    _, _, tt, th = _runs(references, case)
    mine, ref = tt["head.topk_indices"].ravel(), th["head.topk_indices"].ravel()
    k = ref.size
    shared = len(set(mine.tolist()) & set(ref.tolist()))
    strong = ref[th["head.topk_values"].ravel() >= STRONG]
    print(case, f"top-{k} shared {shared}, strong {len(strong)}, positional {np.mean(mine == ref):.4f}")
    assert k == 500 and shared >= 0.99 * k
    assert set(strong.tolist()) <= set(mine.tolist())
    np.testing.assert_allclose(np.sort(tt["head.topk_values"].ravel()), np.sort(th["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, _, th = _runs(references, case)
    tf = TapRegistry(include=SELECTION_TAPS)
    rt.network(b.prepared, tf, topk_indices=th["head.topk_indices"])
    assert np.array_equal(tf["head.topk_indices"], th["head.topk_indices"])
    assert np.array_equal(tf["label_pred"], th["label_pred"])
    worst = {}
    for name in SELECTION_TAPS:
        if name in ("head.topk_indices", "label_pred"):
            continue
        x, y = tf[name], th[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"], th["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


@pytest.mark.parametrize("case", ["rosbag_0144", "pandaset_019_f40", "ns0103_k20"])
def test_reproduces_the_research_goldens(references, case):
    """The research reference's goldens of the same frame (rosbag: ``golden_leveled/golden_lidar.npz``; public:
    ``public_data/golden/<id>_lidar.npz``): identical voxel tensors, the BEV (when stored), the sigmoid heatmap,
    the top-500 set and the Autoware detections (labels after the remapper, centre within 1 cm, score within 1e-4)."""
    a, _, tt, _ = _runs(references, case)
    g = _data.research_golden("ps019_f40" if case == "pandaset_019_f40" else case)
    prep = a.prepared
    assert np.array_equal(prep.coors, g["coors"]) and np.array_equal(prep.voxel_set.num_points,
                                                                     g["num_points_per_voxel"])
    assert np.array_equal(prep.vfe, g["vfe"])
    if "bev_lidar" in g.files:
        assert pcc(tt["sparse.bev"], g["bev_lidar"]) >= PCC_MIN
    if "heatmap" in g.files:
        assert pcc(tt["head.heat"], g["heatmap"]) >= PCC_MIN
    if "topk_indices" in g.files:
        shared = len(set(tt["head.topk_indices"].ravel().tolist()) & set(g["topk_indices"].ravel().tolist()))
        assert shared >= 495
    det = g["detections"]                                # label (Autoware), score, x, y, z, l, w, h, yaw, vx, vy
    out = a.output
    m = match_detections(out.boxes, out.label_ids, out.scores, det[:, [2, 3, 4, 5, 6, 7, 8]], det[:, 0].astype(int),
                         det[:, 1], 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