# SPDX-License-Identifier: Apache-2.0
"""The host half of ``api.METEOR`` (no device, no weights): the identity, the runtime params (the C++ renderer's knobs,
``host.postprocess.PostConfig``), ``_prepare`` on the shipped sample (the golden feed, per-stream state) and
``_postprocess`` on synthetic outputs (the same ``host.result`` code as the CPU reference). The shipped sample is
the synthetic eight-camera frame; the PandaSet 019 frame 40 sample is staged (``tests/paths.py``).

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

import json
from pathlib import Path

import numpy as np
import pytest

from tt_meteor import io as tio
from tt_meteor.api import KNOBS, METEOR, Output, load_sample
from tt_meteor.host.outputs import MeteorOutput
from tt_meteor.host.postprocess import PostConfig
from tt_meteor.reference import config as C
from tt_meteor.tests import paths as P

PKG = Path(__file__).resolve().parents[1]


@pytest.fixture
def model():
    m = METEOR.__new__(METEOR)          # the host half only: no weights, no device
    m.compile_params = {}
    m._build_host()
    return m


def test_identity_and_params():
    assert METEOR.MODEL_NAME == "meteor-p150" and METEOR.DEFAULT_REVISION == C.HF_REVISION
    assert list(METEOR.CAMERA_ORDER) == list(C.CAMERAS) and METEOR.LABELS == ["VEHICLE", "VRU"]
    assert Output is MeteorOutput and METEOR.EXTRA_INPUTS == ("ego_speed",)
    cfg = PostConfig()
    for name, (_, _, _, default) in METEOR.RUNTIME_PARAMS.items():
        if name == "det2d_hide":
            assert tuple(int(v) for v in default.split(",")) == cfg.det2d_hide
        elif name != "heads":
            assert getattr(cfg, name) == default, name
    p = METEOR.validate_params({"vehicle_threshold": "0.5", "det2d_hide": "7,9"})
    assert p["vehicle_threshold"] == 0.5 and p["det2d_hide"] == "7,9"
    with pytest.raises(tio.InputError):
        METEOR.validate_params({"bev_nms_iou": 1.5})
    assert KNOBS.defaults().as_dict() == {"INPUT_NORM": "onnx", "DEPTH_MEAN_BINS": "log", "MAX_STREAMS": 16,
                                         "IMAGE_PRECISION": "terms3"}


def test_prepare_the_sample_and_streams(model):
    gold = P.pandaset_golden("inputs")
    kw = load_sample(P.PANDASET_SAMPLE)
    prep = model._prepare(**kw)
    assert prep["frame"].sha256() == gold["input_sha256"] and prep["stream_id"] == "pandaset_019"
    st = prep["state"]
    st.frames = 3
    assert model._prepare(**kw)["state"] is st and st.frames == 3                       # same stream id
    assert model._prepare(**dict(kw, stream={"id": "pandaset_019", "reset": True}))["state"].frames == 0
    with pytest.raises(tio.InputError):
        model._prepare(points=np.zeros((3, 4), np.float32), **kw)
    with pytest.raises(tio.InputError):
        model._prepare(**dict(kw, ego_speed=None))
    assert model.cfg.input_norm == "onnx" and model.cfg.depth_mean_bins == "log"


def test_prepare_the_shipped_sample(model):
    """The shipped synthetic sample: eight present cameras, its preset, the feed sha256 stored at generation."""
    gold = P.shipped_inputs()
    kw = load_sample(P.SHIPPED_SAMPLE)
    assert kw["calibration"]["cameras"].keys() == set(C.CAMERAS)
    prep = model._prepare(**kw)
    assert prep["frame"].sha256() == gold["input_sha256"] and prep["stream_id"] == P.SHIPPED_NAME
    assert prep["frame"].present.all() and gold["present"] == [1] * 8


def test_postprocess_synthetic(model):
    kw = load_sample(P.SHIPPED_SAMPLE)
    prep = model._prepare(**kw)
    rng = np.random.default_rng(0)
    raw = {"hm": np.full((1, 2, 400, 250), -9.0, np.float32), "reg": np.zeros((1, 6, 400, 250), np.float32),
           "lane": rng.integers(0, 3, (1, 800, 500)).astype(np.uint8), "ego": np.zeros((1, 42), np.float32),
           "tl": np.array([[3.0, 0.0, 0.0, 0.0]], np.float32), "stationary": rng.normal(0, 1, (1, 1, 400, 250)),
           "traj": np.zeros((1, 39, 400, 250), np.float32), "lane_logit": rng.normal(0, 1, (1, 9, 800, 500)),
           "depth_mean": np.full((1, 8, 54, 96), 5.0, np.float16)}
    raw["hm"][0, 0, 150, 125] = 3.0
    for si, s in enumerate(C.DET2D_STRIDES):
        h, w = C.FEAT_H * 4 // s, C.FEAT_W * 4 // s
        raw[f"hm2d_s{si}"] = np.full((1, 8, 10, h, w), -20.0, np.float32)
        raw[f"reg2d_s{si}"] = np.zeros((1, 8, 4, h, w), np.float32)
    out = model._postprocess(raw, prep, METEOR.validate_params({"seg_fuse": False}))
    d = out.to_dict()
    assert d["num_detections"] == 1 and d["detections"][0]["center"][0] == pytest.approx(80 - 150 * 0.4)
    assert d["traffic_light"]["state"] == "none" and d["meta"]["stream_id"] == P.SHIPPED_NAME
    assert prep["state"].frames == 1 and prep["state"].mode.prev == 0
    out2 = model._postprocess(raw, prep, METEOR.validate_params({}))                    # seg fusion with the pose
    assert prep["state"].seg.acc is not None and out2.lane.shape == (800, 500)


def test_stored_reference_bodies_agree():
    """``samples/<stem>.reference.json`` (the smoke test's reference) is the CPU reference's /predict body of the
    shipped request, consistent with the decoded golden of tests/goldens, and compares clean with itself."""
    from tt_meteor.server import smoke_test

    gold = P.pandaset_golden("reference")
    body = json.loads(P.PANDASET_REFERENCE.read_text())
    assert body["model"] == "meteor-p150" and body["num_detections"] == len(gold["boxes3d"]) == 8
    assert [round(d["score"], 5) for d in body["detections"]] == [round(b["score"], 5) for b in gold["boxes3d"]]
    assert body["plan"]["mode"] == gold["plan_mode"] and len(body["trajectory"]) == 6
    assert gold["input_sha256"] == P.pandaset_golden("inputs")["input_sha256"]
    metrics, fails = smoke_test.ttaw_smoke.compare_with_reference(body, body, gates=smoke_test.REFERENCE_GATES)
    assert fails == [] and metrics["detections"]["recall"] == 1.0 and metrics["lane"]["agreement"] == 1.0
    assert set(smoke_test.REQUIRED_KEYS) <= set(body)
    assert smoke_test.output_gates(body, "VEHICLE:4,VRU:4") == []


def test_shipped_reference_body():
    """The shipped sample's stored reference (the container smoke's): the documented fields, clean against itself,
    and the smoke's default ``--expect`` holds on it."""
    from tt_meteor.server import smoke_test

    body = json.loads(P.SHIPPED_REFERENCE.read_text())
    assert body["model"] == "meteor-p150" and body["frame_id"] == "base_link" and body["timing_ms"] == {}
    assert set(smoke_test.REQUIRED_KEYS) <= set(body) and len(body["trajectory"]) == 6
    assert body["meta"]["present"] == [True] * 8 and body["meta"]["v0"] == pytest.approx(8.0)
    metrics, fails = smoke_test.ttaw_smoke.compare_with_reference(body, body, gates=smoke_test.REFERENCE_GATES)
    assert fails == [] and metrics["lane"]["agreement"] == 1.0 and "trajectory" in metrics
    assert "{{" not in smoke_test.DEFAULT_EXPECT
    assert smoke_test.output_gates(body, smoke_test.DEFAULT_EXPECT) == []
    assert smoke_test.DEFAULT_INPUT == P.SHIPPED_SAMPLE