#!/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] [--variant seg_det|seg]

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 coloured by the predicted class, with the 3D boxes; tt_ptv3.viz). The default input is the shipped
sample code/tt_ptv3/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_ptv3" / "samples" / "test_pcd.npz"))
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
ap.add_argument("--variant", default="seg_det", choices=("seg_det", "seg"))
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)

# --- the model card snippet --------------------------------------------------------------------------------------
from tt_ptv3 import PTv3

with PTv3.from_pretrained(device_id=args.device_id, variant=args.variant) as model:   # weights -> HF cache, traces captured
    out = model(args.input)                     # base_link cloud: path (.npz/.npy/.pcd/.bin), (N, 4) array, or a PointCloud

print({k: v for k, v in out.class_counts().items() if v})              # input points per class (255 = outside the crop)
for d in (out.detections.to_dicts() if out.detections is not None else [])[: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_ptv3.viz import load_xyz, render_bev  # noqa: E402

dets = out.detections.to_dicts() if out.detections is not None else []
timing = {k: round(v, 1) for k, v in out.timing_ms.items()}
img = render_bev(load_xyz(args.input), out.label_ids, dets,
                 title=f"PTv3 (Autoware ptv3 v4.0) on one Blackhole p150: {Path(args.input).name}",
                 lines=[f"{len(out.label_ids)} points, {len(dets)} boxes; timing_ms {timing}",
                        "Bird's-eye view, forward up, rings every 10 m; points coloured by the predicted class."])
img.save(out_dir / "quickstart_bev.png")
print(f"-> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'}  timing_ms={timing}")