# SPDX-License-Identifier: Apache-2.0
"""The CPU fp32 reference (``tt_streampetr.reference``) against ONNX Runtime on the three deployed ONNX graphs
(PLAN.md 0.3 item 1; ORT runs them as-is, S:6).

Gates:
- before any data-dependent selection: PCC >= 0.9999 (``img_feats``, ``pos_embed`` / ``cone``, every head output
  computed before the top-256: cls, bbox, ``outs_dec``, the decoder inputs);
- after it (the top-256 selection, the post-update memory rows, the detections): set-based or teacher-forced -- the
  top-256 overlap, the post memory with ORT's own selection (recovered from its ``post_memory_embedding``), and the
  detection F1 of the end-to-end runs (same label, 1 m).

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

    cd bundles/streampetr-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_streampetr/tests/test_reference_cpu.py

Inputs: the staged PandaSet sample (frames 40-43 through the full node pipeline) and the research
teacher-forcing records of the public sequences (``research/streampetr/public_data``; nuScenes is internal only).
"""
from __future__ import annotations

import os

import numpy as np
import pytest

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

from tt_streampetr.host import inputs as hin  # noqa: E402
from tt_streampetr.host import postprocess as hp  # noqa: E402
from tt_streampetr.host.agreement import detection_f1, recover_topk_indices, topk_overlap  # noqa: E402
from tt_streampetr.host.temporal import StreamState  # noqa: E402
from tt_streampetr.reference.config import default_config  # noqa: E402
from tt_streampetr.tests import paths  # noqa: E402
from tt_streampetr.ttaw.metrics import pcc  # noqa: E402

REF_PCC = 0.9999                     # PLAN.md 0.3 item 1
CFG = default_config()
THREADS = int(os.environ.get("STREAMPETR_TEST_THREADS", "3"))   # workspace rule: <= 4 CPU threads
PRE_SELECTION = ("all_cls_scores", "all_bbox_preds", "outs_dec", "reference_points", "tgt", "query_pos",
                 "temp_memory", "temp_pos", "query_pos_in")
POST_MEMORY = ("post_memory_embedding", "post_memory_reference_point", "post_memory_egopose", "post_memory_velo")
FRAMES = (40, 41, 42, 43)


@pytest.fixture(scope="module")
def weights():
    w = paths.weights_dir()
    if w is None:
        pytest.skip("camera_streampetr v1.0 weights not found (set STREAMPETR_WEIGHTS_DIR)")
    return w


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

    return OrtOracle(weights, threads=THREADS)


@pytest.fixture(scope="module")
def reference(weights):
    from tt_streampetr.reference.pipeline import StreamPETRReference

    return StreamPETRReference(weights, CFG, threads=THREADS, preprocess_workers=2)


@pytest.fixture(scope="module")
def sample_run(reference, oracle):
    """The staged PandaSet sample (019 frames 40-43, one stream from a cold start) through the reference pipeline
    and, with the same host inputs and the same host memory loop, through ORT."""
    from tt_streampetr.ttaw.golden import TapRegistry

    paths.need_pandaset()
    reference.streams.clear()
    st_ort = StreamState(CFG)
    runs = []
    for f in FRAMES:
        kw = hin.load_sample(paths.pandaset_sample(f))
        taps = TapRegistry(include=("img_feats",))
        det, outs, plan, img = reference(**kw, taps=taps, return_outputs=True)
        frame = hin.prepare_frame(**kw, cfg=CFG)
        p_ort = st_ort.begin(frame.anchor_stamp, frame.stream.T_world_from_ego, frame.camera_stamps)
        pe_o = oracle.position_embedding(**frame.rig.pe_inputs())
        feats_o = oracle.backbone(img)
        feeds = {"x": feats_o, "pos_embed": pe_o["pos_embed"], "cone": pe_o["cone"], "data_ego_pose": p_ort.E,
                 "data_ego_pose_inv": p_ort.Einv, **p_ort.pre}
        outs_o = oracle.head(feeds)
        st_ort.end(p_ort, outs_o)
        runs.append({"frame": f, "img": img, "rig": frame.rig, "plan": plan, "outs": outs, "det": det,
                     "feats_o": feats_o, "pe_o": pe_o, "feeds_o": feeds, "outs_o": outs_o, "plan_o": p_ort})
    return runs


def test_backbone_vs_ort(sample_run):
    for r in sample_run[:2]:
        p = pcc(r["outs"]["img_feats"], r["feats_o"])
        print(f"  frame {r['frame']} img_feats pcc={p:.12f} max_abs="
              f"{np.abs(r['outs']['img_feats'] - r['feats_o']).max():.2e} (|x| <= {np.abs(r['feats_o']).max():.0f})")
        assert p >= REF_PCC


def test_position_embedding_vs_ort(reference, sample_run):
    r = sample_run[0]
    pos, cone = reference.position_embedding(r["rig"])
    p = pcc(pos, r["pe_o"]["pos_embed"])
    print(f"  pos_embed pcc={p:.12f} max_abs={np.abs(pos - r['pe_o']['pos_embed']).max():.2e}; cone max_abs="
          f"{np.abs(cone - r['pe_o']['cone']).max():.2e}")
    assert p >= REF_PCC and pcc(cone, r["pe_o"]["cone"]) >= REF_PCC


def test_head_vs_ort_on_the_same_inputs(reference, sample_run):
    """The head with ORT's own inputs (ORT backbone / PE / memory) at every frame of the sample."""
    for r in sample_run:
        fo = r["feeds_o"]
        pre = {k: fo[k] for k in fo if k.startswith("pre_")}
        out = reference.head(fo["x"], fo["pos_embed"], fo["cone"], fo["data_ego_pose"], fo["data_ego_pose_inv"], pre)
        for k in PRE_SELECTION:
            p = pcc(out[k], r["outs_o"][k])
            assert p >= REF_PCC, (r["frame"], k, p)
        idx_o = recover_topk_indices(r["outs_o"]["post_memory_embedding"], r["outs_o"]["outs_dec"])
        ov = topk_overlap(out["topk_idx"].reshape(-1), idx_o)
        tf = reference.head(fo["x"], fo["pos_embed"], fo["cone"], fo["data_ego_pose"], fo["data_ego_pose_inv"], pre,
                            topk_idx=idx_o)
        worst = min(pcc(tf[k], r["outs_o"][k]) for k in POST_MEMORY)
        print(f"  frame {r['frame']}: cls pcc={pcc(out['all_cls_scores'], r['outs_o']['all_cls_scores']):.10f} "
              f"top-256 overlap={ov:.4f} teacher-forced post-memory min pcc={worst:.10f}")
        assert ov >= 0.99 and worst >= REF_PCC
        assert np.array_equal(tf["post_memory_timestamp"], r["outs_o"]["post_memory_timestamp"])


def test_end_to_end_detections_vs_ort(sample_run):
    """Free-running over the 4 frames (each side with its own memory): the node's objects agree (set-based)."""
    for r in sample_run:
        ours = hp.autoware_postprocess(r["outs"]["all_cls_scores"], r["outs"]["all_bbox_preds"], CFG)["objects"]
        theirs = hp.autoware_postprocess(r["outs_o"]["all_cls_scores"], r["outs_o"]["all_bbox_preds"], CFG)["objects"]
        m = detection_f1(ours, theirs, dist=1.0)
        print(f"  frame {r['frame']}: {len(ours)} vs {len(theirs)} objects, F1={m['f1']:.4f} "
              f"centre err max {m['centre_err_max']:.4f} m, score err max {m['score_err_max']:.2e}")
        assert m["f1"] >= 0.99 and m["centre_err_max"] < 0.05
        assert r["det"].to_dict()["num_detections"] == len(ours)


def _tf_records():
    out = []
    for run, f in (("ns0103", 1), ("ns0916", 0)):
        g = paths.public_golden(run)
        if g is not None and (g / f"tf_{f:04d}.npz").is_file():
            out.append((run, f))
    return out


@pytest.mark.parametrize("run,frame", _tf_records(), ids=[f"{r}-{f}" for r, f in _tf_records()])
def test_head_vs_ort_on_nuscenes_records(reference, oracle, run, frame):
    """A second rig (1600x900 nuScenes cameras, internal validation data): PE and head vs ORT."""
    g = paths.public_golden(run)
    rec, pe = np.load(g / f"tf_{frame:04d}.npz"), np.load(g / "pe_inputs_outputs.npz")
    pe_o = oracle.position_embedding(pe["img_metas_pad"], pe["intrinsics"], pe["img2lidar"])
    pos, cone = reference.net.pe(pe["img_metas_pad"], pe["intrinsics"], pe["img2lidar"])
    assert pcc(pos.detach().numpy(), pe_o["pos_embed"]) >= REF_PCC
    pre = {k: rec[k] for k in rec.files if k.startswith("pre_")}
    feeds = {"x": rec["img_feats"], "pos_embed": pe_o["pos_embed"], "cone": pe_o["cone"],
             "data_ego_pose": rec["data_ego_pose"], "data_ego_pose_inv": rec["data_ego_pose_inv"], **pre}
    outs_o = oracle.head(feeds)
    out = reference.head(rec["img_feats"], pe_o["pos_embed"], pe_o["cone"], rec["data_ego_pose"],
                         rec["data_ego_pose_inv"], pre)
    for k in PRE_SELECTION:
        assert pcc(out[k], outs_o[k]) >= REF_PCC, k