"""
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/

eMac split-model "prepared input equivalence": SplitInputState must reproduce,
byte-exact, the queue semantics of compile_split_runtime's execute_bundle —
the real tinygrad reference graph run on CPU with stub vision/policy runners,
over a multi-frame random sequence with desire rising edges.
"""
from __future__ import annotations

import os

import numpy as np
import pytest

os.environ.setdefault("DEV", "CPU")

from iqpilot.selfdrive.iqmodeld.emac_input_state import EmacInputState, SplitInputState

N_FRAMES_TEST = 30
FRAME_SKIP = 4
IMG_SHAPE = (1, 12, 16, 32)  # small spatial dims: queue math is shape-generic
FB_SHAPE = (1, 25, 512)
DP_SHAPE = (1, 25, 8)
VISION_OUT_LEN = 1576
HIDDEN_SLICE = slice(1064, 1576)

VISION_SHAPES = {"img": IMG_SHAPE, "big_img": IMG_SHAPE}
POLICY_SHAPES = {"desire_pulse": DP_SHAPE, "traffic_convention": (1, 2), "features_buffer": FB_SHAPE}


class _StubRunner:
  """Stands in for OnnxRunner inside execute_bundle: returns a preset output
  and records the materialized inputs it was fed."""

  def __init__(self, out_len: int):
    self.out_len = out_len
    self.next_output: np.ndarray | None = None
    self.captured: dict[str, np.ndarray] | None = None

  def __call__(self, inputs):
    from tinygrad import Tensor
    self.captured = {k: v.numpy().copy() for k, v in inputs.items()}
    out = self.next_output if self.next_output is not None else np.zeros((1, self.out_len), dtype=np.float32)
    return {"outputs": Tensor(out.astype(np.float32))}


@pytest.fixture(scope="module")
def reference():
  from tinygrad import Tensor
  from iqpilot.selfdrive.iqmodeld.tools.compile_split_runtime import _role_executor

  meta_by_role = {
    "vision": {"input_shapes": dict(VISION_SHAPES), "output_slices": {"hidden_state": HIDDEN_SLICE}},
    "policy": {"input_shapes": dict(POLICY_SHAPES), "output_slices": {}},
  }
  vision, policy = _StubRunner(VISION_OUT_LEN), _StubRunner(1000)
  execute_bundle = _role_executor({"vision": vision, "policy": policy}, meta_by_role, FRAME_SKIP)

  feat_q = Tensor(np.zeros((FRAME_SKIP * (FB_SHAPE[1] - 1) + 1, FB_SHAPE[0], FB_SHAPE[2]), dtype=np.float32),
                  device="CPU").contiguous().realize()
  desire_q = Tensor(np.zeros((FRAME_SKIP * DP_SHAPE[1], DP_SHAPE[0], DP_SHAPE[2]), dtype=np.float32),
                    device="CPU").contiguous().realize()
  return execute_bundle, feat_q, desire_q, vision, policy


def test_split_inputs_match_tinygrad_reference(reference):
  from tinygrad import Tensor

  execute_bundle, feat_q, desire_q, vision_stub, policy_stub = reference
  rng = np.random.default_rng(4321)
  state = SplitInputState(FRAME_SKIP, IMG_SHAPE, FB_SHAPE, DP_SHAPE)
  ref_prev_desire = np.zeros(DP_SHAPE[2], dtype=np.float32)

  for frame in range(N_FRAMES_TEST):
    warped = rng.integers(0, 256, (2, 6, IMG_SHAPE[2], IMG_SHAPE[3]), dtype=np.int64).astype(np.uint8)
    raw_desire = np.zeros(DP_SHAPE[2], dtype=np.float32)
    if frame % 3:
      raw_desire[int(rng.integers(0, DP_SHAPE[2]))] = 1.0
    traffic = rng.standard_normal((1, 2)).astype(np.float32)
    vision_out = rng.standard_normal((1, VISION_OUT_LEN)).astype(np.float32)
    vision_stub.next_output = vision_out

    # --- ours ---
    vis_inputs = state.materialize_vision(warped, raw_desire)
    pol_inputs = state.materialize_policy(vision_out[0, HIDDEN_SLICE], traffic[0])

    # --- reference graph: rising edge happens outside execute_bundle (run_fused) ---
    cur = raw_desire.copy()
    cur[0] = 0
    ref_pulse = np.where(cur - ref_prev_desire > 0.99, cur, 0).astype(np.float32)
    ref_prev_desire[:] = cur

    execute_bundle(
      img=Tensor(vis_inputs["img"], device="CPU").realize(),
      big_img=Tensor(vis_inputs["big_img"], device="CPU").realize(),
      feat_q=feat_q, desire_q=desire_q,
      desire=Tensor(ref_pulse, device="CPU").realize(),
      traffic_convention=Tensor(traffic, device="CPU").realize(),
      action_t=Tensor(np.zeros((1, 2), dtype=np.float32), device="CPU").realize(),
    )
    ref = policy_stub.captured
    assert ref is not None

    assert ref["features_buffer"].tobytes() == pol_inputs["features_buffer"].tobytes(), f"features frame {frame}"
    assert ref["desire_pulse"].tobytes() == pol_inputs["desire_pulse"].tobytes(), f"desire frame {frame}"
    assert ref["traffic_convention"].tobytes() == pol_inputs["traffic_convention"].tobytes()
    # vision saw exactly what our img queues materialized
    vref = vision_stub.captured
    assert vref["img"].tobytes() == vis_inputs["img"].tobytes(), f"img frame {frame}"
    assert vref["big_img"].tobytes() == vis_inputs["big_img"].tobytes(), f"big_img frame {frame}"


def test_split_img_queue_matches_fused_state():
  # img/desire mechanics are shared with the fused mirror: same warps must
  # materialize identical img/big_img in both states
  rng = np.random.default_rng(7)
  fused_spec = {
    "img": (IMG_SHAPE, "uint8"), "big_img": (IMG_SHAPE, "uint8"),
    "desire_pulse": (DP_SHAPE, "float32"), "traffic_convention": ((1, 2), "float32"),
    "features_buffer": ((1, 24, 512), "float32"), "action_t": ((1, 2), "float32"),
  }
  fused = EmacInputState(FRAME_SKIP, fused_spec)
  split = SplitInputState(FRAME_SKIP, IMG_SHAPE, FB_SHAPE, DP_SHAPE)
  for _ in range(12):
    warped = rng.integers(0, 256, (2, 6, IMG_SHAPE[2], IMG_SHAPE[3]), dtype=np.int64).astype(np.uint8)
    desire = np.zeros(DP_SHAPE[2], dtype=np.float32)
    f = fused.push_and_materialize(warped, desire, np.zeros(2, dtype=np.float32), np.zeros(2, dtype=np.float32))
    s = split.materialize_vision(warped, desire)
    assert f["img"].tobytes() == s["img"].tobytes()
    assert f["big_img"].tobytes() == s["big_img"].tobytes()