# SPDX-License-Identifier: Apache-2.0
"""The host half of ``api.StreamPETR`` (``_build_host`` / ``_prepare`` / ``_postprocess``; no device, no weights)
against the CPU reference: the shipped sample prepares to the research goldens' inputs, the stream logic follows the
node, and the reference head outputs post-process to the stored reference bodies (the smoke test's
``samples/synthetic_5cam.reference.json`` of the shipped sample, and the staged PandaSet frame 40's).

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

import hashlib
import json

import numpy as np
import pytest

from tt_streampetr import io as tio
from tt_streampetr.api import KNOBS, StreamPETR
from tt_streampetr.host import inputs as hin
from tt_streampetr.tests import paths

SHIPPED_HEAD = paths.GOLDENS_SMALL / "synthetic_5cam_head.npz"
SHIPPED_BODY = paths.SAMPLES / "synthetic_5cam.reference.json"


def _host_model(**params):
    m = object.__new__(StreamPETR)          # no from_pretrained: no weights, no device
    m.compile_params = dict(params)
    m._build_host()
    return m


@pytest.fixture
def model():
    return _host_model(HOST_PRE="numpy")     # the float32 NCHW ``img`` (the tests below hash it)


@pytest.mark.parametrize("which", ["shipped", "pandaset40"])
def test_host_pre_numba_upload_is_bit_identical(which):
    """HOST_PRE=numba (``host.fast``) writes exactly the bf16 rows the numpy path uploads (``img_rows`` + RNE)."""
    from tt_streampetr.host import fast as hfast
    from tt_streampetr.tt.model import img_rows
    from tt_streampetr.ttaw.tensors import float32_to_bf16_bits

    if not hfast.available():
        pytest.skip("numba not installed")
    if which == "shipped":
        manifest = paths.SHIPPED_SAMPLE
    else:
        paths.need_pandaset()
        manifest = paths.pandaset_sample(40)
    ref = _host_model(HOST_PRE="numpy")._prepare(**hin.load_sample(manifest))
    fast = _host_model(HOST_PRE="numba")
    got = fast._prepare(**hin.load_sample(manifest))
    assert "img" not in got and got["img_bits"].dtype == np.uint16
    want = float32_to_bf16_bits(img_rows(ref["img"]))
    assert np.array_equal(got["img_bits"].reshape(want.shape), want)
    again = fast._prepare(**hin.load_sample(manifest))           # the reused buffer
    assert again["img_bits"] is got["img_bits"] and np.array_equal(again["img_bits"].reshape(want.shape), want)


def test_identity_and_params():
    assert StreamPETR.CAMERA_ORDER == ("CAM_FRONT", "CAM_FRONT_LEFT", "CAM_BACK_LEFT", "CAM_FRONT_RIGHT",
                                       "CAM_BACK_RIGHT")
    assert StreamPETR.LABELS == ["CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN"] and StreamPETR.requires_calibration()
    assert StreamPETR.validate_params({}) == {"score_threshold": None, "iou_nms_threshold": 0.5,
                                              "iou_nms_search_distance": 10.0, "circle_nms_distance": 0.0,
                                              "decode_layers": 6, "awml_decode": False, "max_detections": 500}
    for bad in ({"normalization": "awml_training"}, {"decode_layers": 0}, {"awml_decode": 1}):
        with pytest.raises(tio.InputError):
            StreamPETR.validate_params(bad)                       # load-time knob / out of range / not a bool
    assert KNOBS.read(env={}).as_dict() == {"NORMALIZATION": "autoware_main", "USE_TEMPORAL": True,
                                            "MAX_STREAMS": 1, "AUTOWARE_COMPAT": True, "PREPROCESS_FMA": False,
                                            "PREPROCESS_WORKERS": 4, "HOST_PRE": "numba", "HOST_POST": "numba",
                                            "STEM_IN": "kin", "CONV_L1": "s3r45n", "CONV_L1_ABH": "256,128,64,32",
                                            "ESE": "kernel", "ESE_MEAN": "kernel", "OSA_CAT": "kmm", "CONV_CFG": "adb,wdb", "CONV_CHAIN": "kcat45", "RESHAPE": "kernel", "KMM_BUF": "auto", "OSA_MM": "rowblock", "KV_EPI": True,
                                            "SELF_KV": "buffer", "S26_KERNEL": True, "SEL": "mk",
                                            "PRECISION": "HiFi4+fp32+l1acc:w=fp32",
                                            "ATTN_EXP": "exact", "XATTN": "x8", "ATTN_FP32_ACC": True, "HEAD_ACT": "float32", "EXACT_MAP": "gather",
                                            "KEY_SPLIT": True, "TEMPORAL_EXACT": True}


def test_serve_env_pins_the_knob_defaults():
    """``tt-model.yaml`` ``serve.env`` pins every load-time knob at its default."""
    yaml = pytest.importorskip("yaml")
    path = paths.BUNDLE_DIR / "tt-model.yaml"
    if not path.is_file():
        pytest.skip("no tt-model.yaml (installed package)")
    env = yaml.safe_load(path.read_text())["serve"]["env"]
    pins = KNOBS.serve_env()
    assert {k: str(env.get(k)) for k in pins} == pins


def test_prepare_sample_and_stream_semantics(model):
    paths.need_pandaset()
    kw = hin.load_sample(paths.pandaset_sample(40))
    prep = model._prepare(**kw)
    assert prep["img"].shape == (1, 5, 3, 480, 640) and prep["img"].dtype == np.float32
    plan = prep["plan"]
    assert plan.fresh and plan.ts == 0 and plan.skipped is None and prep["stream_id"] == "pandaset_019"
    gold = paths.public_golden("ps019_f40_43_cold")
    if gold is not None:                                       # bit-identical to the research golden inputs
        want = {r["frame"]: r["img_f32_sha256"] for r in json.loads((gold / "inputs_sha256.json").read_text())}[40]
        got = {c: hashlib.sha256(prep["img"][0, i].tobytes()).hexdigest() for i, c in enumerate(StreamPETR.CAMERA_ORDER)}
        assert got == want
        con = np.load(paths.ship_sample() / "contract.npz")
        assert np.array_equal(plan.E, con["data_ego_pose"][0]) and np.array_equal(prep["frame"].rig.intrinsics,
                                                                                  con["intrinsics"])
    # the next frame continues the stream (the host part of the memory loop runs in _forward with the device)
    st = prep["state"]
    st.end(plan, {"post_memory_timestamp": np.zeros((1, 1280, 1)), "post_memory_embedding": np.zeros((1, 1280, 256)),
                  "post_memory_reference_point": np.zeros((1, 1280, 3)), "post_memory_egopose": np.zeros((1, 1280, 4, 4)),
                  "post_memory_velo": np.zeros((1, 1280, 2))})
    p2 = model._prepare(**hin.load_sample(paths.pandaset_sample(41)))["plan"]
    assert not p2.fresh and abs(float(p2.ts) - 0.099967) < 1e-5 and p2.frame_index == 2
    # frame 40 again: its manifest resets the stream
    assert model._prepare(**kw)["plan"].fresh
    # a single frame without a stream; a point cloud is refused; a stream without its time stamp is refused
    single = dict(kw)
    single.pop("stream")
    assert model._prepare(**single)["stream_id"] is None
    with pytest.raises(tio.InputError):
        model._prepare(points=np.zeros((3, 4), np.float32))
    with pytest.raises(tio.InputError):
        model._prepare(**dict(kw, stream={"id": "x"}))


def test_camera_sync_failure_skips_the_frame(model):
    paths.need_pandaset()
    kw = hin.load_sample(paths.pandaset_sample(40))
    model._prepare(**dict(kw, stream=dict(kw["stream"], id="s")))
    late = hin.load_sample(paths.pandaset_sample(41))
    late["images"][3]["timestamp_s"] += 0.3                      # CAM_FRONT_RIGHT 0.3 s late (> 0.15 s)
    prep = model._prepare(**dict(late, stream=dict(late["stream"], id="s")))
    assert prep["plan"].skipped == "camera_sync" and "img" not in prep
    out = model._postprocess(model._forward(prep), prep, StreamPETR.validate_params({}))
    assert len(out) == 0 and out.meta["skipped"] == "camera_sync"


def _check_postprocess(model, manifest, head, body_path):
    prep = model._prepare(**hin.load_sample(manifest))
    with np.load(head) as g:
        raw = {"all_cls_scores": g["all_cls_scores"], "all_bbox_preds": g["all_bbox_preds"]}
    body = model._postprocess(raw, prep, StreamPETR.validate_params({})).to_dict()
    ref = json.loads(body_path.read_text())
    assert body["detections"] == ref["detections"] and body["num_detections"] == ref["num_detections"]
    assert body["meta"]["decoder_layer"] == ref["meta"]["decoder_layer"]
    return ref


def test_shipped_sample_postprocess_matches_its_stored_reference(model):
    """The shipped synthetic sample: the reference head outputs post-process to the stored smoke reference, which
    holds one object per class with large margins (code/scripts/make_sample.py)."""
    assert paths.SHIPPED_SAMPLE.is_file() and SHIPPED_HEAD.is_file() and SHIPPED_BODY.is_file()
    ref = _check_postprocess(model, paths.SHIPPED_SAMPLE, SHIPPED_HEAD, SHIPPED_BODY)
    assert sorted(d["label"] for d in ref["detections"]) == sorted(StreamPETR.LABELS)
    prep = model._prepare(**hin.load_sample(paths.SHIPPED_SAMPLE))
    assert prep["plan"].fresh and prep["img"].shape == (1, 5, 3, 480, 640)


def test_pandaset_postprocess_matches_the_stored_reference(model):
    paths.need_pandaset()
    _check_postprocess(model, paths.pandaset_sample(40), paths.PANDASET_GOLDENS / "pandaset_019_f40_head.npz",
                       paths.PANDASET_SAMPLES / "pandaset_019_f40.reference.json")


def test_stream_frame_without_ego_pose_is_skipped(model):
    """The node skips a cycle whose map <- base_link TF lookup fails and keeps its state (PORT_LOG section 6 Q2)."""
    kw = hin.load_sample(paths.SHIPPED_SAMPLE)
    first = model._prepare(**dict(kw, stream=dict(kw["stream"], id="pose")))
    first["state"].end_device(first["plan"])
    frames = first["state"].frames
    nopose = {k: v for k, v in kw["stream"].items() if k != "T_world_from_ego"}
    prep = model._prepare(**dict(kw, stream=dict(nopose, id="pose", reset=False, timestamp_s=nopose["timestamp_s"] + 0.1)))
    assert prep["plan"].skipped == "no_ego_pose" and "img" not in prep and prep["state"].frames == frames
    out = model._postprocess(model._forward(prep), prep, StreamPETR.validate_params({}))
    assert len(out) == 0 and out.meta["skipped"] == "no_ego_pose"
    single = dict(kw)
    single.pop("stream")                                        # a stateless frame needs no pose
    assert model._prepare(**single)["plan"].skipped is None