#!/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 manifest] [--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 result: the BEV lane map, the 3D boxes with their futures and the three ego paths, the selected one in
green). The default input, the shipped synthetic sample, is found relative to this file (runs from any directory);
a manifest given on the command line is relative to the current directory.
"""
import argparse
import json
import math
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parents[1]
ap = argparse.ArgumentParser()
ap.add_argument("input", nargs="?", default=str(REPO / "code" / "tt_meteor" / "samples" / "synthetic_8cam.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_meteor import METEOR, load_sample

with METEOR.from_pretrained(device_id=args.device_id) as model:   # weights -> HF cache, trace captured
    out = model(**load_sample(args.input))                          # 8 cameras + calibration + ego speed + stream

for d in out.to_dicts():
    print(d["label"], round(d["score"], 3), d["center"], d["size"], round(d["yaw"], 3), "stationary" if d["stationary"] else "")
body = out.to_dict()
print("plan: mode", body["plan"]["mode"], "path", [[round(v, 2) for v in p] for p in body["trajectory"]],
      "| traffic light:", body["traffic_light"]["state"], "| 2D boxes:", sum(len(v) for v in body["detections_2d"].values()))
# ------------------------------------------------------------------------------------------------------------------

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

# bird's-eye view: x forward (up), y left (left); 0.2 m lane cells, 50 m ahead, 25 m behind, +-25 m to the sides
from PIL import Image, ImageDraw  # noqa: E402

PALETTE = np.array([(0, 0, 0), (90, 90, 90), (140, 90, 160), (0, 200, 200), (255, 255, 255), (255, 40, 40),
                    (255, 140, 0), (240, 220, 60), (40, 60, 140)], np.uint8)   # METEOR's viz palette (lane classes)
lane = np.asarray(out.lane, np.uint8)[150:525, 125:375]                         # x +50 .. -25 m, y +25 .. -25 m
S = 2                                                                            # px per 0.2 m cell
img = Image.fromarray(PALETTE[lane]).resize((lane.shape[1] * S, lane.shape[0] * S), Image.NEAREST)
draw = ImageDraw.Draw(img)


def px(x, y):
    return (25.0 - y) * 5 * S, (50.0 - x) * 5 * S


for d in body["detections"]:
    (x, y), (L, W), yaw = d["center"], d["size"], d["yaw"]
    c, s = math.cos(yaw), math.sin(yaw)
    corners = [px(x + c * a - s * b, y + s * a + c * b) for a, b in ((L / 2, W / 2), (L / 2, -W / 2),
                                                                     (-L / 2, -W / 2), (-L / 2, W / 2))]
    col = (160, 160, 160) if d["stationary"] else ((255, 215, 0) if d["label_id"] == 0 else (255, 0, 255))
    draw.polygon(corners, outline=col, width=2)
    draw.line([px(x, y), px(x + c * L / 2, y + s * L / 2)], fill=col, width=2)
    if d.get("future") and not d["stationary"]:
        draw.line([px(x, y)] + [px(a, b) for a, b in d["future"]], fill=col, width=1)
for k, path in enumerate(body["plan"]["paths"]):
    sel = k == body["plan"]["mode"]
    draw.line([px(0, 0)] + [px(a, b) for a, b in path], fill=(0, 255, 0) if sel else (0, 150, 60), width=3 if sel else 1)
ex, ey = px(0, 0)
draw.polygon([(ex, ey - 9), (ex - 6, ey + 7), (ex + 6, ey + 7)], fill=(255, 255, 255))
img.save(out_dir / "quickstart_bev.png")
print(f"{out} -> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'}  "
      f"timing_ms={ {k: round(v, 1) for k, v in out.timing_ms.items()} }")