# SPDX-License-Identifier: Apache-2.0
"""The fp32 CPU reference (``tt_meteor.reference``) against ONNX Runtime on the deployed ``meteor_v157c3Z.onnx``
(PLAN.md 0.3 item 1: PCC >= 0.9999 before any data-dependent selection, set-based after it).

Runs in the research venv (onnxruntime + torch + onnx); skipped where onnxruntime, torch or the ONNX is missing::

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

1. The shipped sample (PandaSet 019 frame 40) through ORT with every reference tap as an extra graph output and through
   the reference: all 48 internal taps and the 19 outputs. Float tensors PCC >= 0.9999 (measured 1.0000000000);
   boolean / argmax maps agreement >= 0.9999 (``lift.valid`` exact). ``depth_mean`` follows an argmax (a +-2-bin
   window around the softmax maximum): it is compared set-based, >= 99.9 % of the pixels within 2 fp16 ulps.
2. The stored ORT goldens of three more frames from two more rigs (METEOR's in-domain valday #40, nuScenes scene-0103
   key-frame 9, PandaSet 090 frame 40; research/meteor and research/meteor/public_data), fed exactly as ORT saw them;
   valday also against its fp16 internal taps.
3. ``input_norm="imagenet"`` (D12 option) against ORT on the graph with the training normalisation inserted.
4. The TT port's exact rewrites (``reference.port_form``) teacher-forced from the ORT taps; the host lift geometry
   (``host.calib``) bit-exact against the graph's; the host decode (``host.result``) against the research decode
   (``research/meteor/public_data/scripts/me_public_common.py``).
A JSON report with every number is written to ``logs/meteor/test_reference_cpu.json`` of the workspace.
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import numpy as np
import pytest

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

from tt_meteor.host.calib import lift_geometry, lift_ground_points  # noqa: E402
from tt_meteor.host.postprocess import PostConfig  # noqa: E402
from tt_meteor.host.preprocess import MeteorFrame, load_scene_frame  # noqa: E402
from tt_meteor.host.result import build_output  # noqa: E402
from tt_meteor.reference import config as C  # noqa: E402
from tt_meteor.reference import port_form as PF  # noqa: E402
from tt_meteor.reference.model import OUTPUT_NAMES, MeteorNet  # noqa: E402
from tt_meteor.reference.ort import ORT_TAPS, OrtOracle  # noqa: E402
from tt_meteor.reference.pipeline import MeteorReference  # noqa: E402
from tt_meteor.reference.weights import MeteorWeights, find_onnx  # noqa: E402
from tt_meteor.ttaw.golden import TapRegistry  # noqa: E402
from tt_meteor.ttaw.metrics import pcc  # noqa: E402

PKG = Path(__file__).resolve().parents[1]
from tt_meteor.tests import paths as P  # noqa: E402

SAMPLE = P.PANDASET_SCENE / "pandaset_019_f40"            # the staged PandaSet sample (tests/paths.py)
ROOT = Path("/home/ubuntu/experiments/tt-models")
RESEARCH = ROOT / "research" / "meteor"
THREADS = int(os.environ.get("METEOR_TEST_THREADS", "4"))
PCC_MIN = 0.9999
AGREE_MIN = 0.9999
REPORT: dict = {}

try:
    ONNX = find_onnx()
except FileNotFoundError:
    pytest.skip("meteor_v157c3Z.onnx not found (METEOR_WEIGHTS_DIR or assets/meteor/hf_meteor)",
                allow_module_level=True)


def _report(section: str, rows: dict) -> None:
    REPORT.setdefault(section, {}).update(rows)
    out = ROOT / "logs" / "meteor"
    if out.is_dir():
        (out / "test_reference_cpu.json").write_text(json.dumps(REPORT, indent=1, sort_keys=True))


def _cmp(name: str, test, ref) -> dict:
    t, r = np.asarray(test), np.asarray(ref)
    assert t.shape == r.shape, f"{name}: shape {t.shape} vs {r.shape}"
    if r.dtype in (np.bool_, np.uint8) or t.dtype in (np.bool_, np.uint8):
        return {"agreement": float((t == r).mean())}
    tf, rf = t.astype(np.float64), r.astype(np.float64)
    return {"pcc": pcc(tf, rf), "max_abs": float(np.abs(tf - rf).max()), "ref_max": float(np.abs(rf).max())}


def _depth_mean_setwise(test, ref) -> dict:
    t, r = np.asarray(test, np.float64), np.asarray(ref, np.float64)
    ulp = 2.0 ** (np.floor(np.log2(np.maximum(np.abs(r), 2.0 ** -14))) - 10)       # fp16 spacing at the value
    return {"within_2ulp": float((np.abs(t - r) <= 2 * ulp).mean()), "pcc": pcc(t, r)}


def _check(name: str, m: dict) -> None:
    if "agreement" in m:
        assert m["agreement"] >= AGREE_MIN, f"{name}: agreement {m['agreement']:.6f}"
    else:
        assert m["pcc"] >= PCC_MIN, f"{name}: pcc {m['pcc']:.8f}"


# ------------------------------------------------------------------------------------------------- fixtures

@pytest.fixture(scope="module")
def weights():
    return MeteorWeights(ONNX)


@pytest.fixture(scope="module")
def reference(weights):
    return MeteorReference(weights=weights, threads=THREADS)


@pytest.fixture(scope="module")
def frame():
    P.need_pandaset()
    return load_scene_frame(SAMPLE, 0)


@pytest.fixture(scope="module")
def ort_taps(frame):
    oracle = OrtOracle(ONNX, taps=list(ORT_TAPS), threads=THREADS)
    out = oracle(frame.feed())
    del oracle
    return out


@pytest.fixture(scope="module")
def ref_taps(reference, frame):
    taps = TapRegistry()
    reference.run_frame(frame, taps)
    return taps.to_dict()


# ---------------------------------------------------------------------------------------------------- tests

def test_graph_check(weights):
    ops = weights.check_graph()
    assert ops["Conv"] == 190 and ops["Gemm"] == 15 and ops.get("BatchNormalization", 0) == 0
    assert len(weights.conv_modules()) == 190
    assert weights.paint_classes() == C.PAINT_CLASSES
    assert abs(weights.risk_gate() - 2.2981195) < 1e-6 and abs(weights.dec_gate() - 0.0336921) < 1e-6
    np.testing.assert_array_equal(lift_ground_points(), weights.lift_ground_points()[0])   # torch linspace


def test_every_tap_matches_ort(ort_taps, ref_taps):
    rows = {}
    for name in list(ORT_TAPS) + list(OUTPUT_NAMES):
        if name == "depth_mean":
            rows[name] = _depth_mean_setwise(ref_taps[name], ort_taps[name])
            assert rows[name]["within_2ulp"] >= 0.999 and rows[name]["pcc"] >= PCC_MIN, rows[name]
            continue
        rows[name] = _cmp(name, ref_taps[name], ort_taps[name])
        _check(name, rows[name])
    assert rows["lift.valid"]["agreement"] == 1.0 and rows["lift.grid"]["max_abs"] == 0.0
    _report("pandaset_019_f40 vs ORT (live, all taps)", rows)
    print("\n".join(f"  {k:22s} {v}" for k, v in rows.items()))


GOLDENS = [
    ("valday_f040", RESEARCH / "golden" / "valday_f040_ort.npz",
     RESEARCH / "golden" / "valday_f040_intermediates_fp16.npz"),
    ("ns0103_f0009", RESEARCH / "public_data" / "golden" / "ns0103" / "full_f0009.npz", None),
    ("ps090_f0020", RESEARCH / "public_data" / "golden" / "ps090" / "full_f0020.npz", None),
]
# research intermediate name -> reference tap name (research/meteor/scripts/dump_intermediates.py)
FP16_TAPS = {"f_fpn": "img.fpn", "depth_logits": "depth.logits", "ctx_painted": "ctx.painted",
             "lift_weight": "lift.weight", "lift_bev_04": "lift.bev", "raw_bev": "bev.raw", "fused_bev": "bev.fused",
             "lane_logits_pre_refiner": "lane.pre", "det_feat": "det.feat", "hm_pre_refiner": "det.hm_pre",
             "reg_pre_refiner": "det.reg_pre", "occ_feat": "occ.feat", "traj_feat256": "traj.feat",
             "ego_mlp": "ego.mlp", "ego_plus_attn": "ego.plus_attn", "ego_plus_sem": "ego.plus_sem",
             "ego_plus_kin": "ego.plus_kin", "ego_after_risk": "ego.after_risk", "ego_after_dec": "ego.after_dec"}


@pytest.mark.parametrize("name,path,taps_path", GOLDENS, ids=[g[0] for g in GOLDENS])
def test_stored_ort_goldens(reference, name, path, taps_path):
    if not path.is_file():
        pytest.skip(f"{path} not present (research goldens are not shipped)")
    g = np.load(path)
    present = np.array([bool(np.asarray(g["in_imgs"])[0, i].any()) for i in range(C.N_CAMS)])
    frame = MeteorFrame(g["in_imgs"], g["in_K"], g["in_T_cam_ego"], g["in_v0"], present)
    want_taps = taps_path is not None and taps_path.is_file()
    taps = TapRegistry(include=list(FP16_TAPS.values())) if want_taps else TapRegistry(enabled=False)
    out = reference.run_frame(frame, taps)
    rows = {}
    for k in OUTPUT_NAMES:
        if k == "depth_mean":
            rows[k] = _depth_mean_setwise(out[k], g[k])
            assert rows[k]["within_2ulp"] >= 0.999 and rows[k]["pcc"] >= PCC_MIN, (k, rows[k])
            continue
        rows[k] = _cmp(k, out[k], g[k])
        _check(f"{name}.{k}", rows[k])
    if want_taps:
        it = np.load(taps_path)
        for rk, tk in FP16_TAPS.items():
            rows["tap." + tk] = _cmp(tk, taps[tk], it[rk].astype(np.float32))
            assert rows["tap." + tk]["pcc"] >= PCC_MIN, (tk, rows["tap." + tk])
    _report(f"{name} vs stored ORT golden", rows)


def test_imagenet_variant(weights, frame):
    """The oracle graph normalises every camera; the port also zeroes an absent camera after normalising (the
    trained "absent" input, SPEC section 2), so the parity check runs with every camera marked present."""
    oracle = OrtOracle(ONNX, threads=THREADS, input_norm="imagenet")
    want = oracle(frame.feed())
    del oracle
    ref = MeteorReference(weights=weights, threads=THREADS, cfg=C.default_config(input_norm="imagenet"))
    all_present = MeteorFrame(frame.imgs, frame.K, frame.T_cam_ego, frame.v0, np.ones(C.N_CAMS, bool))
    out = ref.run_frame(all_present)
    rows = {}
    for k in OUTPUT_NAMES:
        if k == "depth_mean":
            rows[k] = _depth_mean_setwise(out[k], want[k])
            assert rows[k]["within_2ulp"] >= 0.999, rows[k]
            continue
        rows[k] = _cmp(k, out[k], want[k])
        _check(f"imagenet.{k}", rows[k])
    _report("pandaset_019_f40 input_norm=imagenet vs ORT (normalisation inserted after /Div)", rows)


def test_port_form_rewrites(weights, ort_taps, frame):
    """Every rewrite of ``reference.port_form`` reproduces the literal graph (teacher-forced from ORT)."""
    F = torch.nn.functional
    net = MeteorNet(weights, check_graph=False)
    T = lambda k: torch.from_numpy(np.ascontiguousarray(ort_taps[k]))
    rows = {}
    w, b, off = PF.stem_input_fold(weights, "onnx")
    x = torch.from_numpy(frame.imgs.reshape(C.N_CAMS, 3, C.IMG_H, C.IMG_W).astype(np.float32))
    y = F.max_pool2d(F.relu(F.conv2d(x - torch.from_numpy(off).reshape(1, 3, 1, 1), torch.from_numpy(w),
                                     torch.from_numpy(b), stride=2, padding=3)), 3, 2, 1)
    rows["stem_input_fold"] = _cmp("stem", y, ort_taps["img.stem"])
    wm, bg, bt = PF.tgate_tfuse_merged(weights)
    gate, t0 = PF.fused_from_merged(F.conv2d(T("bev.raw"), torch.from_numpy(wm)), bg, bt)
    rows["tgate_tfuse_merged.gate"] = _cmp("gate", gate, ort_taps["bev.gate"])
    rows["tgate_tfuse_merged.fused"] = _cmp("fused", T("bev.raw") + net.block(t0, "tfuse3_slim/tfuse3_slim.3"),
                                            ort_taps["bev.fused"])
    wd, bd = PF.merged_heads(weights)["det"]
    o = F.conv2d(T("det.feat"), torch.from_numpy(wd), torch.from_numpy(bd))
    rows["merged_heads.hm"] = _cmp("hm", o[:, :2], ort_taps["det.hm_pre"])
    rows["merged_heads.reg"] = _cmp("reg", o[:, 2:], ort_taps["det.reg_pre"])
    w2, b2 = PF.merged_heads(weights)["det2d_s0"]
    det2d_stem = net.block(T("img.fpn"), "det2d_stem")
    o2 = F.conv2d(det2d_stem, torch.from_numpy(w2), torch.from_numpy(b2))
    rows["merged_heads.hm2d_s0"] = _cmp("hm2d", o2[:, :10].reshape(1, 8, 10, 108, 192), ort_taps["hm2d_s0"])
    wp, bp = PF.paint_conv21(weights)
    p = torch.softmax(net.nan_to_num(T("seg2d.logits")), dim=1)
    rows["paint_conv21"] = _cmp("ctx", net.conv(T("img.fpn"), "ctx") + F.conv2d(p, torch.from_numpy(wp),
                                                                                torch.from_numpy(bp)),
                                ort_taps["ctx.painted"])
    wl, bl = PF.traj_stem_live(weights)
    s = net.block(F.relu(F.conv2d(T("bev.fused"), torch.from_numpy(wl), torch.from_numpy(bl), stride=2, padding=1)),
                  "traj_stem/traj_stem.3")
    rows["traj_stem_live"] = _cmp("traj.stem", s, ort_taps["traj.stem"])
    geom = lift_geometry(frame.K, frame.T_cam_ego, points=weights.lift_ground_points()[0])
    rows["lift_from_tables"] = _cmp("lift.bev", PF.lift_from_tables(T("ctx.painted"), T("depth.prob"), geom),
                                    ort_taps["lift.bev"])
    for k, m in rows.items():
        assert m["pcc"] >= 0.99999, (k, m)
        assert m["max_abs"] <= 1e-4 * max(1.0, m["ref_max"]), (k, m)
    _report("port_form rewrites (teacher-forced from ORT)", rows)


def test_host_lift_geometry_bit_exact(ort_taps, frame, weights):
    geom = lift_geometry(frame.K, frame.T_cam_ego, points=weights.lift_ground_points()[0])
    np.testing.assert_array_equal(geom.grid, ort_taps["lift.grid"][:, :, 0, :])
    np.testing.assert_array_equal(geom.valid, ort_taps["lift.valid"])
    st = geom.stats()
    assert st["pairs_total"] == 173111 and st["k_max"] == 3            # research/meteor/public_data README (ps019)
    _report("host lift geometry (pandaset_019)", st)


def test_host_decode_matches_research(ort_taps, frame):
    """``host.result.build_output`` on ORT's outputs == the research decode of METEOR's C++ rules."""
    scripts = RESEARCH / "public_data" / "scripts"
    if not (scripts / "me_public_common.py").is_file():
        pytest.skip("research scripts not present")
    sys.path.insert(0, str(scripts))
    import me_public_common as R

    out = {k: ort_taps[k] for k in OUTPUT_NAMES}
    res = build_output(out, frame, PostConfig(), None)
    rd, healthy = R.decode_3d(out)
    assert healthy == res.stationary_healthy and len(rd) == len(res.boxes3d)
    for a, b in zip(sorted(res.boxes3d, key=lambda z: -z.score), rd):
        assert (a.cls, a.stationary, a.future_mode, tuple(a.cell)) == (b["cls"], b["stationary"], b["future_mode"],
                                                                       tuple(b["cell"]))
        assert abs(a.score - b["score"]) < 1e-6 and abs(a.x - b["x"]) < 1e-4 and abs(a.y - b["y"]) < 1e-4
        assert abs(a.yaw - b["yaw"]) < 1e-5 and np.abs(np.asarray(a.future) - np.asarray(b["future"])).max() < 1e-3
    r2 = R.decode_2d(out)
    for cam in range(C.N_CAMS):
        assert len(res.boxes2d[cam]) == len(r2[cam]), C.CAMERAS[cam]
        for a, b in zip(res.boxes2d[cam], r2[cam]):
            assert a.cls == b["cls"] and abs(a.score - b["score"]) < 1e-6, C.CAMERAS[cam]
            assert max(abs(a.cx - b["cx"]), abs(a.cy - b["cy"]), abs(a.w - b["w"]), abs(a.h - b["h"])) < 1e-3
    assert res.plan["mode"] == R.select_mode(out["ego"])
    _report("host decode vs research decode (pandaset_019, ORT outputs)",
            {"boxes3d": len(rd), "boxes2d": [len(x) for x in r2], "mode": int(res.plan["mode"])})