#!/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.npz] [--out-dir examples/output]

Writes <out-dir>/quickstart.json (the same JSON as POST /predict: the 8 s ego trajectory, the turn-indicator command
and the predicted paths of the neighbours) and <out-dir>/quickstart_bev.png (a bird's-eye view of the input tensors
with the plan: lanes, route lanes, stop lines and road borders, the neighbours with their predicted 8 s paths, the ego
plan with a dot every second). The default input is found relative to this file (runs from any directory); an input
given on the command line is relative to the current directory and must hold the 15 raw planner tensors
(`DiffusionPlanner.INPUT_SCHEMA`).
"""
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_diffusion_planner" / "samples" / "kashiwanoha_dense.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_diffusion_planner import DiffusionPlanner

with DiffusionPlanner.from_pretrained(device_id=args.device_id) as model:   # weights -> your HF cache, traces captured
    out = model(inputs=args.input)   # the 15 raw planner tensors: .npz path, its bytes, or {name: array}

print(out.columns)            # x, y, yaw, cos, sin, velocity, acceleration (base_link, 0.1-8.0 s)
print(out.poses[:5])
print(out.turn_indicator["command_name"], out.predicted_agents.shape)
# ------------------------------------------------------------------------------------------------------------------

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

# A bird's-eye view of what the model saw and what it planned (ego frame: x forward, y left; pillow only).
import numpy as np                                   # noqa: E402
from PIL import Image, ImageDraw                     # noqa: E402

from tt_diffusion_planner import load_inputs         # noqa: E402

raw = {k: v[0] for k, v in load_inputs(args.input).items()}      # the same decoder as model(inputs=...)
plan = out.poses[:, :2]
fwd = max(60.0, float(plan[:, 0].max()) + 15.0)
x0, x1 = -20.0, fwd                                  # forward range (m)
half = (x1 - x0) / 2                                 # lateral half-width (m): a square view
W = H = 800
s = W / (2 * half)                                   # pixels per metre


def px(xy):
    """ego-frame metres [N, 2] -> image pixels (forward up, left to the left)."""
    xy = np.asarray(xy, np.float64).reshape(-1, 2)
    return [(W / 2 - y * s, H - (x - x0) * s) for x, y in xy]


img = Image.new("RGB", (W, H), (252, 252, 251))
d = ImageDraw.Draw(img)
for name, fill in (("lanes", (236, 235, 231)), ("route_lanes", (205, 226, 251))):
    t = raw[name]
    for lane in t[np.abs(t[:, :, :8]).sum(axis=(1, 2)) > 0]:
        left, right = lane[:, :2] + lane[:, 4:6], lane[:, :2] + lane[:, 6:8]
        d.polygon(px(np.concatenate([left, right[::-1]])), fill=fill)
for lane in raw["lanes"][np.abs(raw["lanes"][:, :, :8]).sum(axis=(1, 2)) > 0]:
    for off in (4, 6):
        d.line(px(lane[:, :2] + lane[:, off:off + 2]), fill=(195, 194, 183), width=1)
for ls in raw["line_strings"][np.abs(raw["line_strings"]).sum(axis=(1, 2)) > 0]:
    stop = ls[0, 2] > 0.5                            # line-string type: stop line, else road border
    d.line(px(ls[:, :2]), fill=(11, 11, 11) if stop else (82, 81, 78), width=4 if stop else 2)
colors = {8: (235, 104, 52), 9: (232, 123, 164), 10: (27, 175, 122)}   # vehicle, pedestrian, bicycle
nb = raw["neighbor_agents_past"]
paths = dict(zip(out.meta["predicted_agent_rows"], out.predicted_agents))
for i in np.flatnonzero(np.abs(nb[:, -1, :8]).sum(axis=1) > 0):
    x, y, c, sn, w, length = (float(v) for v in nb[i, -1, [0, 1, 2, 3, 6, 7]])
    col = colors[8 + int(np.argmax(nb[i, -1, 8:11]))]
    if i in paths:
        d.line(px(np.concatenate([[[x, y]], paths[i][:, :2]])), fill=col, width=1)
    f, lt = np.array([c, sn]), np.array([-sn, c])
    corners = [np.array([x, y]) + a * max(length, 0.5) / 2 * f + b * max(w, 0.5) / 2 * lt
               for a, b in ((1, 1), (1, -1), (-1, -1), (-1, 1))]
    d.polygon(px(corners), fill=col)
wb, length, w = (float(v) for v in raw["ego_shape"])  # wheel base, length, width; base_link = rear axle centre
r = (length - wb) / 2
d.polygon(px([(wb + r, w / 2), (wb + r, -w / 2), (-r, -w / 2), (-r, w / 2)]), fill=(11, 11, 11))   # ego
d.line(px(np.concatenate([[[0.0, 0.0]], plan])), fill=(42, 120, 214), width=4)
for u, v in px(plan[9::10]):                         # one dot per second
    d.ellipse([u - 5, v - 5, u + 5, v + 5], fill=(42, 120, 214), outline=(252, 252, 251), width=2)
d.text((10, 8), f"{Path(args.input).name}: 8 s ego plan (blue, a dot per second), "
                f"turn indicator {out.turn_indicator['command_name']}", fill=(11, 11, 11))
d.text((10, 24), f"{len(paths)} neighbours with predicted 8 s paths; view {2 * half:.0f} m wide, ego frame "
                 "(forward up)", fill=(82, 81, 78))
img.save(out_dir / "quickstart_bev.png")
print(f"{len(out)} poses, turn {out.turn_indicator['command_name']} -> {out_dir / 'quickstart.json'}, "
      f"{out_dir / 'quickstart_bev.png'}  timing_ms={ {k: round(v, 2) for k, v in out.timing_ms.items()} }")