# SPDX-License-Identifier: Apache-2.0
"""Host tests of the API wiring (no device): identity and runtime parameters, the host half of the hooks
(``_build_host``, ``_prepare``, ``_postprocess``) with the fp32 torch reference network in the device's place, against
the stored CPU reference of the samples (the shipped test.pcd and the staged PandaSet frame) (compared with the container smoke's own comparator,
``ttaw.server.smoke``); the shipped samples' voxel, VFE and rulebook hashes (``tests/goldens/voxels.json``); the
input contract.

    TT_VISIBLE_DEVICES=none python -m pytest -q -p fake_ttnn_plugin code/tt_bevfusion/tests/test_api_host.py
"""
from __future__ import annotations

import hashlib
import json

import numpy as np
import pytest

from tt_bevfusion import io as tio
from tt_bevfusion.api import BEVFusion
from tt_bevfusion.host.postprocess import RUNTIME_PARAMS
from tt_bevfusion.reference.config import NODE_PARAMS
from tt_bevfusion.tests import _data
from tt_bevfusion.ttaw.decode import AUTOWARE_LABELS
from tt_bevfusion.ttaw.server import smoke

SAMPLES = ("pandaset_019_f40", "test_pcd")


def _sha(a) -> str:
    return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()


def test_identity_and_runtime_params():
    assert tuple(BEVFusion.VARIANTS) == ("lidar",) and BEVFusion.DEFAULT_VARIANT == "lidar"
    assert tuple(BEVFusion.LABELS) == AUTOWARE_LABELS and BEVFusion.LABELS[7] == "PEDESTRIAN"
    assert BEVFusion.RUNTIME_PARAMS is RUNTIME_PARAMS and BEVFusion.EXTRA_INPUTS == ()
    assert BEVFusion.DEFAULT_REVISION == "e1bf164b909d3c6c2642012f9b5ee4d827753220"
    assert BEVFusion.ALLOW_PATTERNS[0] == "bevfusion_lidar.onnx"
    p = BEVFusion.validate_params({})
    assert p["score_threshold"] == 0.1 and p["remap_classes"] is True and p["max_detections"] is None
    for k in ("circle_nms_dist_threshold", "iou_nms_threshold", "iou_nms_search_distance_2d"):
        assert p[k] == NODE_PARAMS[k]                       # Autoware's defaults
    assert BEVFusion.validate_params({"score_threshold": "0.5"})["score_threshold"] == 0.5
    for bad in ({"score_threshold": 1.5}, {"remap_classes": 1}, {"max_detections": 0}, {"iou_nms_threshold": 1.5},
                {"sampling": "clean"}):
        with pytest.raises(tio.InputError):
            BEVFusion.validate_params(bad)


def _host_model() -> BEVFusion:
    m = BEVFusion.__new__(BEVFusion)                        # no device: the host half only
    m.weights_path = _data.weights_dir()
    m.variant = "lidar"
    m._build_host()
    if m.pre.plan_backend == "levels":       # M0 (DEVICE_PLAN=1) builds no host maps: these tests check the host plan
        m.pre.plan_backend = "xcolumns"      # and feed its maps to the CPU reference (test_fast_host covers "levels")
    return m


@pytest.mark.parametrize("sample", SAMPLES)
def test_shipped_sample_voxels_and_rulebooks(sample):
    """The host pre-processing reproduces the voxel tensors, the VFE and every neighbour map of the stored goldens
    bit for bit (the device inputs of M1-M2)."""
    gold = json.loads((_data.GOLDENS / "voxels.json").read_text())[sample]
    path = _data.sample_path(f"{sample}.npz")
    assert hashlib.sha256(path.read_bytes()).hexdigest() == gold["input_sha256"]
    prep = _host_model()._prepare(points=path)
    v = prep.voxel_set
    assert prep.num_voxels == gold["num_voxels"] and _sha(v.coors) == gold["coors_sha256"]
    assert _sha(v.num_points) == gold["num_points_sha256"] and _sha(v.voxels) == gold["voxels_sha256"]
    assert _sha(prep.vfe) == gold["vfe_sha256"]
    assert prep.plan.counts() == gold["levels"] and prep.plan.pairs() == gold["pairs"]
    assert {k: _sha(m) for k, m in prep.plan.maps.items()} == gold["maps_sha256"]
    assert prep.plan.bucket == gold["bucket"]


@pytest.fixture(scope="module")
def net():
    torch = pytest.importorskip("torch")
    pytest.importorskip("onnx")
    from tt_bevfusion.reference.model import BEVFusionTorch
    from tt_bevfusion.reference.weights import load_weights

    torch.set_num_threads(2)
    return BEVFusionTorch(load_weights(_data.weights_dir()), threads=2)


@pytest.mark.parametrize("sample", SAMPLES)
def test_hooks_reproduce_the_stored_reference(net, sample):
    m = _host_model()
    prep = m._prepare(points=_data.sample_path(f"{sample}.npz"))
    raw = net.forward(prep.vfe, prep.plan)
    out = m._postprocess(raw, prep, m.validate_params({}))
    body = out.to_dict()
    ref = json.loads(_data.sample_path(f"{sample}.reference.json").read_text())
    metrics, fails = smoke.compare_with_reference(body, ref, gates={"max_center_dist": 0.01, "min_recall": 1.0,
                                                                    "min_precision": 1.0, "max_score_err": 1e-4})
    assert fails == [], (fails, metrics)
    assert body["meta"]["counts"] == ref["meta"]["counts"] and body["meta"]["variant"] == "lidar"
    assert body["meta"]["num_voxels"] == ref["meta"]["num_voxels"]
    assert len(body["meta"]["orientation_availability"]) == body["num_detections"] == len(body["meta"]["proposal"])


def test_prepare_refuses_a_frame_over_the_largest_bucket(monkeypatch):
    """VERIFY_PORT L2: a frame whose sparse level counts exceed the largest capacity bucket is refused before the
    device runs, with an explicit client error (``io.InputError``, HTTP 400) that names the levels; nothing is
    clipped. Emulated on test.pcd (21,359 voxels) with tiny buckets."""
    from tt_bevfusion import api as api_mod
    from tt_bevfusion.host import preprocess as hp

    tiny = {"s": (1_024, 1_024, 512, 256, 256), "full": (4_096, 4_096, 2_048, 1_024, 1_024)}
    assert hp.select_bucket({"l1": 4_000, "l2": 4_096, "l3": 2_048, "l4": 1_024, "out": 1_024}, tiny) == ("full", False)
    assert hp.select_bucket({"l1": 4_000, "l2": 4_097, "l3": 10, "l4": 10, "out": 10}, tiny) == ("full", True)
    real = hp.build_sparse_plan
    monkeypatch.setattr(hp, "build_sparse_plan", lambda coors, **kw: real(coors, **{**kw, "buckets": tiny}))
    monkeypatch.setattr(api_mod, "CAPACITY_BUCKETS", tiny)
    m = _host_model()
    with pytest.raises(tio.InputError, match=r"largest capacity bucket 'full'.*l1.*21,359 > 4,096"):
        m._prepare(points=_data.SAMPLE_TESTPCD)


def test_prepare_contract():
    m = _host_model()
    with pytest.raises(tio.InputError):
        m._prepare(points=None)
    with pytest.raises(tio.InputError):
        m._prepare(points=np.zeros((3, 3), np.float32), images=[b"x"])
    prep = m._prepare(points=np.array([[500.0, 0.0, 0.0]], np.float32))     # no voxel: Autoware skips the frame
    assert prep.skipped == "no_voxels" and m._forward(prep) is None
    out = m._postprocess(None, prep, m.validate_params({}))
    assert out.to_dict()["num_detections"] == 0


def test_viz_render_bev_and_matching():
    """``viz.render_bev`` (the quickstart / demo renderer) draws the shipped sample with its stored reference, and its
    matching follows the gates' rule (same label, BEV centres within 0.2 m)."""
    pytest.importorskip("PIL")
    from tt_bevfusion import viz

    ref = json.loads(_data.SAMPLE_TESTPCD.with_name("test_pcd.reference.json").read_text())["detections"]
    moved = [dict(d, center=[d["center"][0] + (0.3 if i == 0 else 0.0)] + d["center"][1:]) for i, d in enumerate(ref)]
    assert viz.match_pairs(moved, ref) == ([0], [0])
    assert viz.match_pairs(ref, ref) == ([], [])
    pts = viz.load_xyz(_data.SAMPLE_TESTPCD)
    assert pts.shape == (48_048, 4) and not pts[:, 3].any()
    img = viz.render_bev(pts, moved, reference=ref, title="test", lines=["caption"], width=480)
    assert img.size[0] == 480 and img.size[1] > 480