"""Send a prepared sample to the Alpamayo2-Super server and print the trajectory error.
  $PY scripts/client_example.py data/samples/030c760c_5100000.pt [--url http://localhost:20000] [--no-reasoning] [--samples 6]
The request format (images[camera][frame] base64 PNG, camera_ids, ego_history_xyz [16,3], ego_history_rot [16,3,3])
is what a real client would send; here the frames come from the prepared sample's raw camera images.
"""
import argparse
import base64
import io
import json
import time

import numpy as np
import requests
import torch
from PIL import Image


def png_b64(chw: torch.Tensor) -> str:
    im = Image.fromarray(chw.permute(1, 2, 0).numpy().astype(np.uint8))
    buf = io.BytesIO()
    im.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


ap = argparse.ArgumentParser()
ap.add_argument("sample")
ap.add_argument("--url", default="http://localhost:20000")
ap.add_argument("--no-reasoning", action="store_true")
ap.add_argument("--greedy", action="store_true")
ap.add_argument("--samples", type=int, default=1)
ap.add_argument("--seed", type=int, default=42)
args = ap.parse_args()

smp = torch.load(args.sample, weights_only=False)
frames = smp["image_frames"]  # [cams, 4, 3, H, W] uint8
req = {
    "images": [[png_b64(frames[c, f]) for f in range(frames.shape[1])] for c in range(frames.shape[0])],
    "camera_ids": smp["camera_indices"].tolist(),
    "ego_history_xyz": smp["ego_history_xyz"][0, 0].tolist(),
    "ego_history_rot": smp["ego_history_rot"][0, 0].tolist(),
    "ego_future_xyz": smp["ego_future_xyz"][0, 0].tolist(),
    "reasoning": not args.no_reasoning, "greedy": args.greedy, "seed": args.seed, "num_traj_samples": args.samples,
}
t0 = time.time()
r = requests.post(f"{args.url}/predict", json=req, timeout=600)
wall = time.time() - t0
r.raise_for_status()
out = r.json()
print(json.dumps({k: out[k] for k in ("chain_of_causation", "generated_tokens", "prompt_tokens", "seed", "timing_ms", "ade_m", "minADE_m") if k in out}, indent=1))
print(f"wall {wall*1e3:.0f} ms; first waypoints {np.round(np.array(out['trajectory_xyz'])[:3], 2).tolist()}")