# SPDX-License-Identifier: Apache-2.0
"""The load-time knobs that are not the default serve configuration, on the device (VERIFY_PORT L4; needs one p150):

    bin/devrun -t 1800 -- python -m pytest -q -s code/tt_streampetr/tests/test_knobs_device.py

- ``NORMALIZATION=autoware_0.52`` + ``USE_TEMPORAL=0`` (one model): frames 40 and 41 of the staged PandaSet sample on
  one stream against the fp32 CPU reference with the same two knobs (the node's per-frame memory wipe and the 0.52.1
  colour statistics). Gate ``knobs.f1_min`` (1 m, same label, the worst of the two frames) >= 0.9, the teacher-forced
  F1 level of PLAN 2.9; frozen like the other gates (``test_knobs_device.gates.json``).
- ``MAX_STREAMS=2`` (a second model): two streams interleaved frame by frame give bit-identical detections to the same
  frames sent as one stream alone (the second stream's memory lives in its own DRAM bank, PLAN D16); a third stream
  id evicts the least recently used one, which then restarts fresh.

The default configuration is covered by ``test_e2e_device.py``. The PandaSet sample is staged (git-ignored), so
inside the tt-models workspace a missing sample FAILS (``tests/paths.need_pandaset``).
"""
from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, Dict, List

import pytest

from tt_streampetr.tests import paths
from tt_streampetr.ttaw.golden import GateRegistry

pytestmark = pytest.mark.device

GATES = GateRegistry.for_test(__file__, {"knobs.f1_min": (0.9, "min", "f1")})
LOGS = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models")) / "logs" / "streampetr"
FRAMES = (40, 41)


def _dets(body: Dict[str, Any]) -> List[Dict[str, Any]]:
    return [{"x": d["center"][0], "y": d["center"][1], "label": d["label_id"], "score": d["score"]}
            for d in body["detections"]]


def _weights():
    wd = paths.weights_dir()
    if wd is None:
        pytest.skip("the camera_streampetr v1.0 ONNX files were not found (STREAMPETR_WEIGHTS_DIR or the HF cache)")
    return wd


def _requests(sid: str, frames=FRAMES):
    from tt_streampetr.host.inputs import load_sample

    out = []
    for k, f in enumerate(frames):
        kw = load_sample(paths.pandaset_sample(f))
        kw["stream"] = dict(kw["stream"], id=sid, reset=k == 0)
        out.append(kw)
    return out


def test_normalization_052_and_no_temporal(device):
    """``autoware_0.52`` + ``USE_TEMPORAL=0`` on the device vs the CPU reference with the same knobs."""
    from tt_streampetr import StreamPETR
    from tt_streampetr.host.agreement import detection_f1
    from tt_streampetr.reference.pipeline import StreamPETRReference

    paths.need_pandaset()
    wd = _weights()
    ref = StreamPETRReference(wd, normalization="autoware_0.52", use_temporal=False, threads=4,
                              preprocess_workers=4)
    cpu = [ref(**kw).to_dict() for kw in _requests("knobs-cpu")]
    m = StreamPETR.from_pretrained(device=device, weights_dir=str(wd), normalization="autoware_0.52",
                                   use_temporal=False)
    with m:
        assert str(m.knobs.NORMALIZATION) == "autoware_0.52" and not bool(m.knobs.USE_TEMPORAL)
        tt = [m(**kw).to_dict() for kw in _requests("knobs-tt")]
    rec = {}
    for f, a, b in zip(FRAMES, tt, cpu):
        assert a["meta"]["fresh"] == (f == FRAMES[0]) and a["meta"]["frame_index"] == b["meta"]["frame_index"]
        rec[f] = detection_f1(_dets(a), _dets(b), dist=1.0)
        print(f"autoware_0.52 + no temporal, frame {f}: TT {a['num_detections']} / CPU {b['num_detections']} "
              f"objects; {json.dumps(rec[f])}")
    LOGS.mkdir(parents=True, exist_ok=True)
    (LOGS / "knobs_norm052_notemporal.json").write_text(json.dumps(rec, indent=1))
    r = GATES.require("knobs.f1_min", min(x["f1"] for x in rec.values()))
    print(r)


def test_two_streams_in_banks(device):
    """``MAX_STREAMS=2``: interleaved streams == each stream alone (bit for bit); a third id evicts the LRU one."""
    from tt_streampetr import StreamPETR

    paths.need_pandaset()
    wd = _weights()
    frames = (40, 41, 42)
    m = StreamPETR.from_pretrained(device=device, weights_dir=str(wd), max_streams=2)
    with m:
        assert int(m.knobs.MAX_STREAMS) == 2
        alone = [m(**kw).to_dict()["detections"] for kw in _requests("alone", frames)]
        a_req, b_req = _requests("bank-a", frames), _requests("bank-b", frames)
        a_out, b_out = [], []
        for ka, kb in zip(a_req, b_req):                          # a1 b1 a2 b2 a3 b3: a bank swap before each call
            a_out.append(m(**ka).to_dict())
            b_out.append(m(**kb).to_dict())
        for i, f in enumerate(frames):
            assert a_out[i]["detections"] == alone[i], f"stream a frame {f} differs from the stream alone"
            assert b_out[i]["detections"] == alone[i], f"stream b frame {f} differs from the stream alone"
            assert a_out[i]["meta"]["fresh"] == (i == 0) and b_out[i]["meta"]["fresh"] == (i == 0)
        # a third stream evicts the least recently used one (a); a's next frame then starts fresh
        m(**_requests("bank-c", frames[:1])[0])
        nxt = _requests("bank-a", (43,))[0]
        nxt["stream"]["reset"] = False
        again = m(**nxt).to_dict()
        assert again["meta"]["fresh"] and sorted(m._banks) == ["bank-a", "bank-c"]
        print(f"MAX_STREAMS=2: interleaved == alone for frames {frames} ({[len(x) for x in alone]} objects); "
              f"eviction restarts the LRU stream fresh")