#!/usr/bin/env python3
"""Offline BPU policy smoke test for RDK X5.

This does not connect to the robot. It loads a BPU .bin and one raw float32
input file, runs the C++ DNN API backend repeatedly, and optionally checks the
known RobotLab 00000.bin reference output captured from hrt_model_exec.
"""

import argparse
import time
from pathlib import Path

import numpy as np

from bpu_policy import BpuInferLibPolicy


HERE = Path(__file__).parent.resolve()
DEFAULT_MODEL = (
    HERE.parent / "bpu_quantization" / "mapper_output_6500_gemm" /
    "policy_robotlab_6500_int16_gemm.bin"
)
DEFAULT_INPUT = HERE.parent / "bpu_quantization" / "calibration_data" / "00000.bin"
REFERENCE_00000 = np.array([
    0.537439,
    1.032869,
    0.042010,
    -0.463015,
    0.515008,
    1.349627,
    0.014652,
    -1.518314,
    1.012123,
    -0.346840,
    0.086353,
    -0.755527,
], dtype=np.float32)


def main():
    parser = argparse.ArgumentParser(description="Offline BPU policy smoke test")
    parser.add_argument("--bpu-model", default=str(DEFAULT_MODEL))
    parser.add_argument("--input-bin", default=str(DEFAULT_INPUT))
    parser.add_argument("--repeat", type=int, default=1000)
    parser.add_argument("--check-reference-00000", action="store_true")
    args = parser.parse_args()

    input_path = Path(args.input_bin).expanduser().resolve()
    data = np.fromfile(input_path, dtype=np.float32)
    policy = BpuInferLibPolicy(args.bpu_model)
    if data.size != policy.input_size:
        raise ValueError(
            f"{input_path} has {data.size} float32 values, "
            f"but model expects {policy.input_size}"
        )
    action = policy(data)
    print("action", np.array2string(action, precision=6))
    print("action_max_abs", float(np.max(np.abs(action))))

    if args.check_reference_00000:
        if policy.input_size != 450:
            raise ValueError("--check-reference-00000 is only valid for the 450-dim RobotLab reference input")
        diff = np.abs(action - REFERENCE_00000)
        print("reference_max_abs_diff", float(diff.max()))
        print("reference_mean_abs_diff", float(diff.mean()))
        if diff.max() > 1e-5:
            raise RuntimeError("BPU output does not match the known 00000.bin reference")

    repeats = max(1, int(args.repeat))
    t0 = time.perf_counter()
    for _ in range(repeats):
        policy(data)
    elapsed_ms = (time.perf_counter() - t0) * 1000.0
    print(f"repeat={repeats} avg_ms={elapsed_ms / repeats:.6f}")


if __name__ == "__main__":
    main()
