#!/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 [sample.json] [--out-dir examples/output] [--ego-pose]

Writes <out-dir>/quickstart.json (the same JSON as POST /predict) and <out-dir>/quickstart_bev.png (a bird's-eye view
of the detections). The input is a sample manifest (six camera images + a calibration preset,
``code/tt_bevdet/samples/*.json``); the default, the shipped synthetic test pattern, is found relative to this file
(runs from any directory); an input given on the command line is relative to the current directory. ``--ego-pose``
passes the manifest's ``stream`` (with ``T_world_from_ego``) for the motion-compensated history mode.
"""
import argparse
import json
import math
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
ap = argparse.ArgumentParser()
ap.add_argument("input", nargs="?", default=str(REPO / "code" / "tt_bevdet" / "samples" / "synthetic_6cam.json"))
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
ap.add_argument("--ego-pose", action="store_true", help="use the manifest's stream (ego pose), if it has one")
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)

# --- the model card snippet --------------------------------------------------------------------------------------
from tt_bevdet import BEVDet, load_sample

with BEVDet.from_pretrained(device_id=args.device_id) as model:   # weights -> HF cache, traces captured
    out = model(**load_sample(args.input, with_stream=args.ego_pose))   # six cameras [FL, F, FR, BL, B, BR] + calib

for d in out.to_dicts():
    print(f'{d["label"]:13s} {d["autoware_label"]:10s} {d["score"]:.3f}  centre {d["center"]}  size {d["size"]}  '
          f'yaw {d["yaw"]:+.2f}  v {d["velocity"]}')
# ------------------------------------------------------------------------------------------------------------------

(out_dir / "quickstart.json").write_text(json.dumps(out.to_dict(), indent=1))


def bev_png(dets, path, rng=40.0, size=640):
    """Bird's-eye view of the detections around the ego vehicle (x forward = up, y left = left), 10 m rings."""
    from PIL import Image, ImageDraw

    colours = {"car": (57, 135, 229), "truck": (57, 135, 229), "bus": (57, 135, 229), "trailer": (57, 135, 229),
               "construction_vehicle": (57, 135, 229), "pedestrian": (217, 89, 38), "bicycle": (25, 158, 112),
               "motorcycle": (25, 158, 112)}
    im = Image.new("RGB", (size, size), (26, 26, 25))
    d = ImageDraw.Draw(im, "RGBA")
    s = size / (2 * rng)

    def px(x, y):
        return ((rng - y) * s, (rng - x) * s)

    for r in range(10, int(rng) + 1, 10):
        (u0, v0), (u1, v1) = px(r, r), px(-r, -r)
        d.ellipse([u0, v0, u1, v1], outline=(44, 44, 42), width=1)
    for det in dets:
        x, y, _ = det["center"]
        length, width, _ = det["size"]
        c, sn = math.cos(det["yaw"]), math.sin(det["yaw"])
        q = [px(x + c * a * length - sn * b * width, y + sn * a * length + c * b * width)
             for a, b in ((0.5, 0.5), (0.5, -0.5), (-0.5, -0.5), (-0.5, 0.5))]
        col = colours.get(det["label"], (195, 194, 183))
        d.polygon(q, fill=col + (70,), outline=col)
        d.line([px(x, y), px(x + c * length / 2, y + sn * length / 2)], fill=col, width=2)
        d.text((q[0][0] + 3, q[0][1] - 12), f'{det["label"]} {det["score"]:.2f}', fill=(195, 194, 183))
    ego = [px(a * 2.05, b * 0.9) for a, b in ((1, 1), (1, -1), (-1, -1), (-1, 1))]
    d.polygon(ego, fill=(255, 255, 255))
    d.text((6, size - 16), f"{2 * rng:.0f} m x {2 * rng:.0f} m, rings every 10 m, up = forward (+x)", fill=(137, 135, 129))
    im.save(path)


bev_png(out.to_dicts(), out_dir / "quickstart_bev.png")
print(f"{len(out)} detections -> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'}  "
      f"timing_ms={ {k: round(v, 2) for k, v in out.timing_ms.items()} }")