# SPDX-License-Identifier: Apache-2.0
"""The CPU fp32 reference (``tt_bevdet.reference``) against ONNX Runtime on the deployed ONNX (PLAN.md 0.3 item 1).

ORT cannot run ``bevdet_one_lt_d.onnx`` as a whole (4 TensorRT plugins), so the oracle is the SPEC's (section 6):
ORT on the two plugin-free sub-graphs cut out of the deployed file (``reference.ort``) + the research plugin
re-implementations and goldens. Gates:

- before any data-dependent selection: PCC >= 0.9999 (depth, feat, every head map, the BEV maps);
- after it (decode threshold, top-500, NMS, node filter): set-based detection agreement with the research goldens.

Runs in the research venv (onnxruntime); skipped where onnxruntime, the weights or the inputs are absent::

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

Inputs: the shipped PandaSet sample (always), the research goldens of the nuScenes key-frame ``sample0`` and of the
PandaSet sequence (``research/bevdet``; internal validation data, never shipped).
"""
from __future__ import annotations

import json
import os

import numpy as np
import pytest

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

from tt_bevdet.host import inputs as hin  # noqa: E402
from tt_bevdet.host import postprocess as hp  # noqa: E402
from tt_bevdet.reference import port_form  # noqa: E402
from tt_bevdet.reference.config import HEAD_NAMES, default_config  # noqa: E402
from tt_bevdet.tests import paths  # noqa: E402
from tt_bevdet.ttaw.metrics import match_detections, pcc  # noqa: E402

REF_PCC = 0.9999                     # PLAN.md 0.3 item 1
CFG = default_config()
THREADS = int(os.environ.get("BEVDET_TEST_THREADS", "3"))   # workspace rule: <= 4 CPU threads


@pytest.fixture(scope="module")
def onnx_path():
    p = paths.weights_onnx()
    if p is None:
        pytest.skip("bevdet_one_lt_d.onnx not found (set BEVDET_WEIGHTS_DIR)")
    return p


@pytest.fixture(scope="module")
def oracle(onnx_path):
    from tt_bevdet.reference.ort import OrtOracle

    return OrtOracle(onnx_path, threads=THREADS)


@pytest.fixture(scope="module")
def reference(onnx_path):
    from tt_bevdet.reference.pipeline import BEVDetReference

    return BEVDetReference(onnx_path, threads=THREADS)


@pytest.fixture(scope="module")
def pandaset_f38():
    kw = hin.load_sample(paths.require_pandaset() / "pandaset_019_f38.json")   # staged until PandaSet may ship
    return hin.prepare_frame(kw["images"], kw["calibration"], CFG)


def _report(name, ours, ref):
    p = pcc(ours, ref)
    print(f"  {name:12s} pcc={p:.10f} max_abs={float(np.abs(ours - ref).max()):.3e}")
    return p


# ---------------------------------------------------------------------------------------------- the oracle

def test_subgraphs_are_faithful_extracts(onnx_path, oracle):
    """Every node and initializer of the extracted sub-graphs is the deployed one (no custom op, none changed)."""
    import onnx
    from onnx import numpy_helper

    dep = onnx.load(str(onnx_path))
    nodes = {(n.op_type, tuple(n.input), tuple(n.output)) for n in dep.graph.node}
    inits = {t.name: t for t in dep.graph.initializer}
    for path, n_expect in ((oracle.img_path, 221), (oracle.bev_path, 63)):     # SPEC 6: 221 + 63 nodes
        sub = onnx.load(str(path))
        assert len(sub.graph.node) == n_expect
        assert all((n.op_type, tuple(n.input), tuple(n.output)) in nodes for n in sub.graph.node)
        assert all(n.domain in ("", "ai.onnx") for n in sub.graph.node)
        for t in sub.graph.initializer:
            assert np.array_equal(numpy_helper.to_array(t), numpy_helper.to_array(inits[t.name])), t.name
    with pytest.raises(Exception, match="Preprocess"):
        ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])   # S:410: ORT cannot run it as is


# ---------------------------------------------------------------------------------------------- image encoder

def _encoder_inputs():
    out = [("pandaset_f38", None)]
    if paths.research_golden("sample0.npz") is not None:
        out.append(("sample0", paths.research_golden("sample0.npz")))
    return out


@pytest.mark.parametrize("name,golden", _encoder_inputs(), ids=[n for n, _ in _encoder_inputs()])
def test_image_encoder_vs_ort(reference, oracle, pandaset_f38, name, golden):
    if golden is None:
        crop, calib = pandaset_f38.crop_u8, pandaset_f38.calib
        img, cp, depth, feat = reference.encode(crop, calib)
    else:
        with np.load(golden) as g:
            img, cp = g["input"], g["cam_params"]
        with torch.no_grad():
            d, f = reference.net.image_encoder(img, cp)
        depth, feat = d.numpy(), f.numpy()
    o_depth, o_feat = oracle.encode(img, cp)
    print(f"\n[{name}] image encoder vs ORT")
    assert _report("depth", depth, o_depth) >= REF_PCC
    assert _report("feat", feat, o_feat) >= REF_PCC
    # the TT port's form: host camera gates (BN1d folded, fp64), 1x3 ASPP, GAP bias, softmax over 128
    gates = port_form.camera_gates(reference.weights, cp)
    with torch.no_grad():
        c4, c5 = reference.net.backbone(torch.from_numpy(np.ascontiguousarray(img)))
        x = reference.net.neck(c4, c5)
    p_depth, p_feat = port_form.forward_depth_net_port(reference.net, x, gates, reference.weights, CFG)
    assert _report("depth(port)", p_depth.numpy(), o_depth) >= REF_PCC
    assert _report("feat(port)", p_feat.numpy(), o_feat) >= REF_PCC


# ---------------------------------------------------------------------------------------------- BEV decoder

def _decoder_inputs():
    out = [("random", None)]
    if paths.research_golden("sample0.npz") is not None:
        out.append(("sample0", paths.research_golden("sample0.npz")))
    pub = paths.public_data()
    full = pub / "pandaset" / "goldens" / "port2hz__ps019_2hz" / "f40_full.npz" if pub else None
    if full is not None and full.is_file():
        out.append(("pandaset_f40_motion", full))
    return out


@pytest.mark.parametrize("name,golden", _decoder_inputs(), ids=[n for n, _ in _decoder_inputs()])
def test_bev_decoder_vs_ort(reference, oracle, name, golden):
    if golden is None:
        rng = np.random.default_rng(0)
        bev_in = np.abs(rng.normal(size=(1, 720, 128, 128))).astype(np.float32) * 5
    else:
        with np.load(golden) as g:
            bev_in = g["bev_in"]
    heads = {k: v.numpy() for k, v in reference.net.bev_decoder(bev_in).items()}
    o_heads = oracle.decode(bev_in)
    print(f"\n[{name}] BEV decoder vs ORT")
    for k in HEAD_NAMES:
        assert _report(k, heads[k], o_heads[k]) >= REF_PCC, k


# ---------------------------------------------------------------------------------------------- end to end

def _match(det, golden_boxes, thr=0.2):
    g = [b for b in golden_boxes if b["score"] >= thr]
    gc = np.array([[b["x"], b["y"]] for b in g]).reshape(-1, 2)
    m = match_detections(det.boxes[:, :2], det.label_ids, det.scores, gc, np.array([b["label"] for b in g]),
                         np.array([b["score"] for b in g]), max_dist=0.5)
    return m, len(g)


def test_end_to_end_pandaset_sequence_vs_research(reference):
    """The stateful reference (the API's semantics: fresh stream, identity ego pose, frames 38, 39, 40) against the
    research goldens of the same Autoware-mode run (``public_data/pandaset/goldens/autoware__ps019_f38-40``)."""
    pub = paths.public_data()
    gdir = pub / "pandaset" / "goldens" / "autoware__ps019_f38-40" if pub else None
    if gdir is None or not gdir.is_dir():
        pytest.skip("research PandaSet goldens not present")
    from tt_bevdet.ttaw.golden import TapRegistry

    reference.reset()
    pdir = paths.require_pandaset()
    for n in (38, 39, 40):
        taps = TapRegistry(include=["depth", "feat", "bev.curr", "head.*_0"])
        det = reference(**hin.load_sample(pdir / f"pandaset_019_f{n}.json"), taps=taps)
        with np.load(gdir / f"f{n}.npz") as g:
            print(f"\n[pandaset f{n}] stateful reference vs research ORT golden")
            assert _report("depth", taps["depth"], g["depth"]) >= REF_PCC
            assert _report("feat", taps["feat"], g["feat"]) >= REF_PCC
            assert _report("curr", taps["bev.curr"], g["curr_bevfeat"]) >= REF_PCC
            for k in HEAD_NAMES:
                assert _report(k, taps[f"head.{k}"], g[k]) >= REF_PCC, k
        m, n_gold = _match(det, json.loads((gdir / f"f{n}_boxes.json").read_text()))
        print(f"  boxes>=0.2 ours={len(det)} golden={n_gold} recall={m.recall:.4f} precision={m.precision:.4f}")
        assert m.recall >= 0.99 and m.precision >= 0.99                          # set-based after the selection
        if len(m.score_errors):
            assert float(np.max(np.abs(m.score_errors))) < 1e-3


def test_end_to_end_sample0_reset_frame(reference):
    """nuScenes key-frame sample0 (reset frame, flag 0) from the raw planes: the BEV stages are bit-exact with the
    research golden and the boxes agree with the 52 golden boxes at score >= 0.2 (SPEC 6)."""
    golden = paths.research_golden("sample0.npz")
    if golden is None:
        pytest.skip("research golden sample0.npz not present")
    from tt_bevdet.host import calib as hc
    from tt_bevdet.host import preprocess as hpre

    with np.load(golden) as g:
        planes = g["planes_bgr_u8"]
        crop = hpre.crop_u8(planes, CFG)
        np.testing.assert_array_equal(hpre.normalize(crop, CFG), g["input"])
        cp = g["cam_params"]
        K = np.stack([np.array([[c[0], 0, c[2]], [0, c[1], c[3]], [0, 0, 1]], np.float32) for c in cp[0]])
        R = cp[0][:, [15, 16, 17, 19, 20, 21, 23, 24, 25]].reshape(6, 3, 3)
        t = cp[0][:, [18, 22, 26]]
        calib = hc.calibration_from_arrays(K, R, t)
        tab = reference.tables(calib)
        np.testing.assert_array_equal(tab.ranks_bev, g["ranks_bev"])
        out = reference.run_frame(crop, calib, g["adj_feats"], g["transforms"], int(g["flag"][0, 0]))
        print("\n[sample0] run_frame vs research golden")
        assert _report("curr", out.curr, g["curr_bevfeat"]) >= REF_PCC
        assert _report("bev_in", out.bev_in, g["bev_in"]) >= REF_PCC
        for k in HEAD_NAMES:
            assert _report(k, out.heads[k], g[k]) >= REF_PCC, k
    golden_boxes = json.loads(golden.with_name("sample0_boxes.json").read_text())
    det = hp.boxes_to_detections(out.boxes, CFG)
    m, n_gold = _match(det, golden_boxes)
    print(f"  boxes>=0.2 ours={len(det)} golden={n_gold} recall={m.recall:.4f} precision={m.precision:.4f}")
    assert n_gold == 52 and m.recall >= 0.99 and m.precision >= 0.99