# 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 as Autoware binds it (PLAN.md 0.3 item 1; BUNDLE_CONVENTIONS.md section 10):

- FRNet has no data-dependent selection inside the network (no top-k, no NMS): every module tap
  (``weights.MODULE_TAPS``: encoder point features and voxel max, frustum image, stem, point stem, fusion stem, the
  four stages' blocks / point fusion / pixel fusion / attention / outputs, decoder, head, logits, ``pred_probs``)
  and the opt-in taps (per-voxel max of every point branch, the five up-sampled decoder inputs) at
  **PCC >= 0.9999** against the same tensor of the deployed graph (ORT exposes it as an extra graph output);
- after the selection, i.e. the per-point argmax of the post-processing: **set-based**, the class ids agree on
  >= 99.99 % of the points and every disagreement is a near-tie of the ORT probabilities (top-2 margin < 1e-4);
- cases: the staged PandaSet sample through the request path (both variants), the hardest public frame for bf16 (nuScenes
  0916 kf 20, OT128) and the densest-cell frame (VLS-128 x > 0 half, QT128: up to 60 points per cell);
- ORT itself reproduces the research goldens' ``pred_probs`` bit for bit (deterministic CPU EP), and the TT dataflow
  (``rewrites.py``) is held to the same tap gates against ORT on the OT128 sample and the VLS-128 QT128 frame.

Runs in the research venv (onnxruntime + torch, CPU only; keep threads <= 4 on the shared host)::

    cd bundles/frnet-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_frnet/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_frnet.host.preprocess import build_plan  # noqa: E402
from tt_frnet.reference.model import FRNetTorch  # noqa: E402
from tt_frnet.reference.ort import OrtFRNet  # noqa: E402
from tt_frnet.reference.pipeline import FRNetReference  # noqa: E402
from tt_frnet.reference.rewrites import FRNetPlanTorch, plan_taps_to_literal  # noqa: E402
from tt_frnet.reference.weights import MODULE_TAPS, load_weights  # noqa: E402
from tt_frnet.tests import _data  # noqa: E402
from tt_frnet.ttaw.golden import TapRegistry  # noqa: E402
from tt_frnet.ttaw.metrics import pcc  # noqa: E402

PCC_MIN = 0.9999            # PLAN.md 0.3 item 1
AGREE_MIN = 0.9999          # set-based class ids after the argmax
TIE = 1e-4                  # a disagreement must be an ORT near-tie
THREADS = 4
PATTERNS = list(MODULE_TAPS) + ["backbone.segmax*", "decoder.up*"]


@pytest.fixture(scope="module")
def weights():
    w = _data.weights_dir()
    torch.set_num_threads(THREADS)
    return {v: load_weights(w, v) for v in ("ot128", "qt128")}


def _inputs(case, weights):
    """``(NetworkInputs, research golden pred_probs or None)`` of a case: ``sample.<variant>`` goes through the
    request path (staged PandaSet sample + its calibration), a catalogued frame is read from its research golden."""
    kind, _, variant = case.rpartition(".")
    if kind == "sample":
        _data.need_pandaset()
        ref = FRNetReference(_data.weights_dir(), variant, backend="torch", threads=THREADS, weights=weights[variant])
        prep = ref.pre(_data.SAMPLE, _data.SAMPLE_CALIB)
        assert prep.skipped is None
        _, gold_probs, _ = _data.load_frame(f"{_data.SAMPLE_STEM}.{variant}")
        return prep.inputs, gold_probs
    inputs, gold_probs, _ = _data.load_frame(case)
    return inputs, gold_probs


CASES = ["sample.ot128", "sample.qt128"] + _data.HARD_FRAMES


@pytest.fixture(scope="module", params=CASES)
def run(request, weights):
    """One case's torch and ORT outputs with every tap (module-scoped and parametrized: pytest runs every test of a
    case together and drops the case before the next one; OT128 taps are ~2.5 GB per side)."""
    case = request.param
    variant = case.rpartition(".")[2]
    inputs, gold = _inputs(case, weights)
    tt, to = TapRegistry(include=PATTERNS), TapRegistry(include=PATTERNS)
    probs_t = FRNetTorch(weights[variant])(inputs.points, inputs.coors, inputs.voxel_coors, inputs.inverse_map,
                                           taps=tt)
    ort = OrtFRNet(_data.weights_dir(), variant, threads=THREADS, intermediates=True, tap_patterns=PATTERNS,
                   weights=weights[variant])
    probs_o = ort(inputs.points, inputs.coors, inputs.voxel_coors, inputs.inverse_map, taps=to)
    del ort
    yield case, inputs, gold, probs_t, probs_o, tt, to
    tt.clear()
    to.clear()


def test_taps_match_onnxruntime(run):
    case, inputs, gold, probs_t, probs_o, tt, to = run
    names = to.names()
    assert set(MODULE_TAPS) <= set(names) and {f"backbone.segmax{i}" for i in range(5)} <= set(names)
    worst = {}
    for name in names:
        x, y = tt[name], to[name]
        assert x.shape == y.shape, name
        worst[name] = pcc(x, y)
    print(case, json.dumps({k: round(v, 12) 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
    assert np.abs(probs_t - probs_o).max() < 1e-4


def test_class_ids_set_based_after_the_argmax(run):
    case, inputs, gold, probs_t, probs_o, _, _ = run
    a, b = probs_t.argmax(1), probs_o.argmax(1)
    assert (a == b).mean() >= AGREE_MIN
    flips = np.flatnonzero(a != b)
    if len(flips):
        top2 = np.sort(probs_o[flips], axis=1)[:, -2:]
        assert (top2[:, 1] - top2[:, 0] < TIE).all(), "a disagreement that is not a near-tie"


def test_onnxruntime_reproduces_the_research_golden(run):
    """Deterministic CPU EP: the same inputs give the research golden's pred_probs bit for bit."""
    _, _, gold, _, probs_o, _, _ = run
    assert gold is not None and np.array_equal(probs_o, gold)


def test_tt_dataflow_matches_onnxruntime(run, weights):
    """``rewrites.py`` (sorted + padded points, gather-form scatters, log-step max, K-split, interpolation
    matrices) against the deployed graph at the same tap gates (OT128 sample and the VLS-128 QT128 frame, whose
    cells hold up to 60 points: 6 log-step rounds)."""
    case, inputs, _, _, probs_o, _, to = run
    if case not in ("sample.ot128", "vls128_frame000_fronthalf.qt128"):
        pytest.skip("the TT dataflow is compared on the OT128 sample and the densest-cell frame")
    variant = case.rpartition(".")[2]
    w = weights[variant]
    cfg = FRNetReference(_data.weights_dir(), variant, backend="torch", weights=w).config
    plan = build_plan(inputs, cfg, pre_norm=(w.pre_norm_scale, w.pre_norm_shift))
    tr = TapRegistry(include=PATTERNS)
    out = plan.unsort(FRNetPlanTorch(w)(plan, taps=tr))
    worst = min((pcc(plan_taps_to_literal(n, tr[n], plan), to[n]), n) for n in to.names())
    print(case, "plan vs ORT worst tap", worst, "rounds", plan.stats()["segmax_rounds"])
    assert worst[0] >= PCC_MIN, worst
    assert np.abs(out - probs_o).max() < 1e-4 and (out.argmax(1) == probs_o.argmax(1)).mean() >= AGREE_MIN