#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model card's Python snippet on the shipped sample, on one Blackhole p150.

    pip install -e .                          # once, from the repo root, on top of an environment that has ttnn (tt-metal)
    python examples/quickstart.py [input] [--out-dir examples/output]

Writes <out-dir>/quickstart.json (the same JSON as POST /predict). The input is a sample manifest (the camera images
and their calibration preset, ``samples/*.json`` of the package, read with ``load_sample``); the default is found
relative to this file (runs from any directory), an input given on the command line is relative to the current
directory.
"""
import argparse
import json
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
ap = argparse.ArgumentParser()
ap.add_argument("input", nargs="?", default=str(REPO / "code" / "tt_streampetr" / "samples" / "synthetic_5cam.json"))
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)

# --- the model card snippet --------------------------------------------------------------------------------------
from tt_streampetr import StreamPETR, load_sample

with StreamPETR.from_pretrained(device_id=args.device_id) as model:   # weights -> HF cache, traces captured
    out = model(**load_sample(args.input))                          # five cameras [FRONT, FRONT_LEFT, BACK_LEFT, FRONT_RIGHT, BACK_RIGHT] + calibration + stream={id, timestamp_s, T_world_from_ego}

for d in out.to_dicts()[:10]:
    print(d)
# ------------------------------------------------------------------------------------------------------------------

(out_dir / "quickstart.json").write_text(json.dumps(out.to_dict(), indent=1))
print(f"{len(out)} results -> {out_dir / 'quickstart.json'}  timing_ms={ {k: round(v, 2) for k, v in out.timing_ms.items()} }")