#!/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) and <out-dir>/quickstart_bev.png (a bird's-eye view
of the input points and the detections, tt_transfusion.viz). The default input is the shipped sample
code/tt_transfusion/samples/test_pcd.npz (Autoware's test.pcd, Apache-2.0), found relative to this file, so the script
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_transfusion" / "samples" / "test_pcd.npz"))
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_transfusion import TransFusion

with TransFusion.from_pretrained(device_id=args.device_id) as model:   # weights -> HF cache, trace captured
    out = model(args.input)    # path, (N, C) array, .npy/.npz/.pcd or PointCloud; a time_lag column = pre-densified

for d in out.to_dicts()[:10]:
    print(d["label"], round(d["score"], 3), d["center"], d["size"], round(d["yaw"], 3))
# ------------------------------------------------------------------------------------------------------------------

(out_dir / "quickstart.json").write_text(json.dumps(out.to_dict(), indent=1))
from tt_transfusion.viz import load_xyz, render_bev  # noqa: E402

timing = {k: round(v, 1) for k, v in out.timing_ms.items()}
img = render_bev(load_xyz(args.input), out.to_dicts(),
                 title=f"TransFusion-L on one Blackhole p150: {Path(args.input).name}",
                 lines=[f"{len(out)} detections; timing_ms {timing}",
                        "Bird's-eye view, forward up, rings every 10 m; boxes are Autoware DetectedObjects in base_link."])
img.save(out_dir / "quickstart_bev.png")
print(f"{len(out)} results -> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'}  timing_ms={timing}")