from __future__ import annotations
import argparse, statistics, time
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.ort_dit import OrtDitModule


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=None, help="Local dit_fp16.onnx path")
    ap.add_argument("--runs", type=int, default=50)
    ap.add_argument("--warmup", type=int, default=10)
    ap.add_argument("--no-trt", action="store_true")
    args = ap.parse_args()

    path = args.model or hf_hub_download("patdev/NitroGen-RTX2060-ONNX", "onnx/dit_fp16.onnx")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    mod = OrtDitModule(path, prefer_tensorrt=not args.no_trt).to(device)
    print(f"device={device} provider={mod.provider}")
    if torch.cuda.is_available():
        print("gpu=", torch.cuda.get_device_name(0))

    h = torch.randn(1, 18, 1024, device=device, dtype=torch.float16)
    e = torch.randn(1, 256, 1024, device=device, dtype=torch.float16)
    t = torch.tensor([500], device=device, dtype=torch.int64)

    for _ in range(args.warmup):
        mod(h, e, t)
    if torch.cuda.is_available(): torch.cuda.synchronize()

    times=[]
    for _ in range(args.runs):
        t0=time.perf_counter(); mod(h,e,t)
        if torch.cuda.is_available(): torch.cuda.synchronize()
        times.append((time.perf_counter()-t0)*1000)

    med=statistics.median(times); p95=sorted(times)[max(0,int(len(times)*.95)-1)]
    print(f"DiT median={med:.2f} ms  p95={p95:.2f} ms")
    for steps in (4,8,16):
        print(f"steps={steps:2d}: DiT-only estimate {1000/(med*steps):.2f} action-chunks/s (vision/head overhead excluded)")

if __name__ == "__main__": main()