#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Render the model card's demo media (``media/``) from the TT (p150) outputs. CPU only; research venv (numpy, PIL,
matplotlib, OpenCV); the TT outputs come from a device run (``run_public_tt.py``, ``tests/test_e2e_device.py``).

    PY=/home/ubuntu/experiments/tt-models/tools/research-venv/bin/python
    $PY code/scripts/make_demo.py datasets --ps019 <tt>/ps019/detections.json --ns0103 <tt>/ns0103/detections.json
    $PY code/scripts/make_demo.py tt-vs-cpu --ps019 <tt>/ps019/detections.json --f1 <tt>/ps019/f1.json
    $PY code/scripts/make_demo.py synthetic --tt-body <logs>/e2e_shipped_synthetic_5cam.json

- ``datasets``: camera mosaics (the five network crops with the projected TT boxes + a BEV tile), large BEV plots and a
  sequence GIF of PandaSet 019 and nuScenes scene-0103 through ``research/streampetr/public_data/scripts/
  render_media.py`` (its privacy blur of heads and plates, its attribution footer); the nuScenes files get the
  ``_NC`` (non-commercial) suffix. Never shipped as data: only the renders are copied into ``media/``.
- ``tt-vs-cpu``: PandaSet 019 frame 40 in a bird's-eye view, TT boxes against the CPU fp32 reference (ONNX Runtime)
  boxes, and the per-frame detection F1 of the 80-frame free-running sequence (no camera pixels).
- ``synthetic``: the shipped synthetic sample: its five camera images with the TT boxes projected, and a BEV with the
  TT boxes, the stored CPU reference and the objects the pattern was made for (generated data, Apache-2.0).
"""
from __future__ import annotations

import argparse
import json
import math
import os
import shutil
import subprocess
import sys
from pathlib import Path

import numpy as np

CODE = Path(__file__).resolve().parents[1]
BUNDLE = CODE.parent
MEDIA = BUNDLE / "media"
ROOT = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))
PUBLIC = ROOT / "research" / "streampetr" / "public_data"
RENDER = PUBLIC / "scripts" / "render_media.py"
CLASSES = ["CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN"]
# the reference categorical palette, dark-surface steps, fixed class order (as render_media.py)
CLS_HEX = ["#3987E5", "#D95926", "#199E70", "#C98500", "#D55181"]
SURFACE, INK, INK2, GRID = "#1A1A19", "#FFFFFF", "#C3C2B7", "#45453F"
TT_LABEL = "TT p150 (ETH 12x10, 1 CQ)"


def _run(args):
    print("+", " ".join(str(a) for a in args), flush=True)
    subprocess.run([str(a) for a in args], check=True)


# ------------------------------------------------------------------------------------------------ datasets

def datasets(a) -> int:
    py = sys.executable
    tmp = Path(a.work)
    tmp.mkdir(parents=True, exist_ok=True)
    MEDIA.mkdir(exist_ok=True)
    lab = f"{TT_LABEL}"
    if a.ps019:
        _run([py, "-I", RENDER, "ps019", "--frames", "40", "--bev", "--det", a.ps019, "--source-label", lab,
              "--tag", "tt", "--out", tmp])
        _run([py, "-I", RENDER, "ps019", "--gif", "0:80:2", "--gif-tile", "240", "--anim-formats", "gif",
              "--det", a.ps019, "--source-label", lab, "--tag", "tt", "--out", tmp])
        shutil.copy(tmp / "streampetr_ps019_f0040_mosaic_tt.jpg", MEDIA / "streampetr_pandaset_019_f40_tt.jpg")
        shutil.copy(tmp / "streampetr_ps019_f0040_bev_tt.png", MEDIA / "streampetr_pandaset_019_f40_bev_tt.png")
        shutil.copy(tmp / "streampetr_ps019_seq_tt.gif", MEDIA / "streampetr_pandaset_019_seq_tt.gif")
    if a.ns0103:
        # sweep 117 of the CAM_FRONT sequence = key-frame 20 of scene-0103
        _run([py, "-I", RENDER, "ns0103", "--frames", "117", "--bev", "--det", a.ns0103, "--source-label", lab,
              "--tag", "tt", "--out", tmp])
        shutil.copy(tmp / "streampetr_ns0103_f0117_mosaic_tt.jpg",
                    MEDIA / "streampetr_nuscenes_scene-0103_kf20_tt_NC.jpg")
        shutil.copy(tmp / "streampetr_ns0103_f0117_bev_tt.png",
                    MEDIA / "streampetr_nuscenes_scene-0103_kf20_bev_tt_NC.png")
    for p in sorted(MEDIA.iterdir()):
        print(p.name, p.stat().st_size)
    return 0


# ------------------------------------------------------------------------------------------------ helpers

def _corners(x, y, yaw, l, w):
    c, s = math.cos(yaw), math.sin(yaw)
    pts = np.array([[0.5, 0.5], [0.5, -0.5], [-0.5, -0.5], [-0.5, 0.5], [0.5, 0.5]]) * [l, w]
    return pts @ np.array([[c, s], [-s, c]]) + [x, y]


def _research_rows(dets):
    """research-format dets -> (label, x, y, yaw, l, w, score)."""
    return [(int(d["label"]), d["x"], d["y"], d["yaw"], d["dim_x"], d["dim_y"], d["score"]) for d in dets]


def _body_rows(body):
    return [(int(d["label_id"]), d["center"][0], d["center"][1], d["yaw"], d["size"][0], d["size"][1], d["score"])
            for d in body["detections"]]


def _bev_axes(ax, lim, title):
    ax.set_facecolor(SURFACE)
    ax.set_xlim(lim, -lim)                      # base_link y (left) on a reversed horizontal axis: left = +y
    ax.set_ylim(-lim * 0.45, lim)
    ax.set_aspect("equal")
    for r in range(10, int(lim * 1.5), 10):
        t = np.linspace(0, 2 * np.pi, 200)
        ax.plot(r * np.sin(t), r * np.cos(t), color=GRID, lw=0.6, zorder=0)
        ax.text(-0.3, r + 0.3, f"{r} m", color=INK2, fontsize=7, zorder=1, clip_on=True)
    ax.plot([0], [0], marker=(3, 0, 0), color=INK, ms=9, zorder=5)
    ax.tick_params(colors=INK2, labelsize=7)
    for s in ax.spines.values():
        s.set_color(GRID)
    ax.set_xlabel("base_link y [m] (left)", color=INK2, fontsize=8)
    ax.set_ylabel("base_link x [m] (forward)", color=INK2, fontsize=8)
    ax.set_title(title, color=INK, fontsize=10, loc="left")


def _draw_boxes(ax, rows, solid, lw=1.6, zorder=3):
    for lab, x, y, yaw, l, w, _ in rows:
        c = _corners(x, y, yaw, l, w)
        ax.plot(c[:, 1], c[:, 0], color=CLS_HEX[lab] if solid else INK2, lw=lw if solid else 1.0,
                ls="-" if solid else (0, (3, 2)), zorder=zorder)


def _legend(ax, extra):
    from matplotlib.lines import Line2D

    h = [Line2D([], [], color=CLS_HEX[i], lw=2, label=n) for i, n in enumerate(CLASSES)] + extra
    lg = ax.legend(handles=h, loc="lower left", fontsize=7, facecolor=SURFACE, edgecolor=GRID, labelcolor=INK2,
                   ncol=2)
    lg.set_zorder(10)


# ------------------------------------------------------------------------------------------------ tt vs cpu

def tt_vs_cpu(a) -> int:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    tt = {int(r["frame"]): r["dets"] for r in json.loads(Path(a.ps019).read_text())["results"]}
    cpu = {int(r["frame"]): r["dets"] for r in json.loads((PUBLIC / "golden" / "ps019" / "detections.json")
                                                             .read_text())["results"]}
    f1 = json.loads(Path(a.f1).read_text())
    frames = f1["frames"]
    k = 40
    fig = plt.figure(figsize=(10.4, 6.6), facecolor=SURFACE)
    gs = fig.add_gridspec(2, 2, width_ratios=[1.15, 1], height_ratios=[1, 1], wspace=0.22, hspace=0.42)
    ax = fig.add_subplot(gs[:, 0])
    _bev_axes(ax, 45, f"PandaSet 019 frame 40: TT {len(tt[k])} vs CPU {len(cpu[k])} boxes")
    ax.set_ylim(-25, 50)
    _draw_boxes(ax, _research_rows(cpu[k]), solid=False, zorder=2)
    _draw_boxes(ax, _research_rows(tt[k]), solid=True)
    _legend(ax, [Line2D([], [], color=INK2, lw=1, ls=(0, (3, 2)), label="CPU fp32 reference")])
    fr = [x["frame"] for x in frames]
    for row, (key, name, ref) in enumerate((("f1", "detection F1 vs CPU (1 m, same class)", None),
                                            ("score_err_max", "max |score difference| of matched boxes", None))):
        b = fig.add_subplot(gs[row, 1])
        b.set_facecolor(SURFACE)
        v = [x.get(key, np.nan) for x in frames]
        b.plot(fr, v, color=CLS_HEX[0], lw=2)
        b.set_title(name, color=INK, fontsize=10, loc="left")
        b.tick_params(colors=INK2, labelsize=7)
        for s in b.spines.values():
            s.set_color(GRID)
        b.grid(color=GRID, lw=0.5)
        b.set_xlabel("frame (10 Hz, one stream from frame 0)", color=INK2, fontsize=8)
        if key == "f1":
            b.set_ylim(0.5, 1.02)
            m = float(np.mean(v))
            b.axhline(m, color=INK2, lw=1, ls=(0, (3, 2)))
            b.text(fr[-1], m - 0.05, f"mean {m:.3f}", color=INK2, fontsize=8, ha="right")
    fig.text(0.01, 0.015, "TT = this port on one p150 (ETH 12x10, 1 CQ), free-running (its own memory); CPU = the "
             "fp32 reference of the same network (ONNX Runtime), free-running. Boxes only, no camera pixels.\n"
             "Contains data from PandaSet (Scale AI and Hesai), https://pandaset.org, CC BY 4.0 + PandaSet Dataset "
             "Terms. Changes: model outputs drawn in a bird's-eye view. Not endorsed by Scale AI or Hesai.",
             color=INK2, fontsize=7)
    out = MEDIA / "streampetr_pandaset_019_tt_vs_cpu.png"
    fig.savefig(out, dpi=110, facecolor=SURFACE)
    print("wrote", out, out.stat().st_size)
    return 0


# ------------------------------------------------------------------------------------------------ synthetic

def synthetic(a) -> int:
    import cv2
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D
    from PIL import Image

    sys.path.insert(0, str(CODE))
    from tt_streampetr.server.smoke_test import REFERENCE_GATES  # noqa: F401  (same rule as the smoke)

    samples = CODE / "tt_streampetr" / "samples"
    man = json.loads((samples / "synthetic_5cam.json").read_text())
    cal = json.loads((CODE / "tt_streampetr" / "calib" / "synthetic_5cam.json").read_text())
    ref = json.loads((samples / "synthetic_5cam.reference.json").read_text())
    rec = json.loads(Path(a.tt_body).read_text())
    tt = rec.get("body", rec)
    rows_tt, rows_ref = _body_rows(tt), _body_rows(ref)
    order = ["CAM_FRONT_LEFT", "CAM_FRONT", "CAM_FRONT_RIGHT", "CAM_BACK_LEFT", "CAM_BACK_RIGHT"]
    tiles = {}
    for cam in order:
        img = np.asarray(Image.open(samples / man["images"][cam]).convert("RGB")).copy()
        K = np.asarray(cal["cameras"][cam]["intrinsics"], np.float64)
        Tinv = np.linalg.inv(np.asarray(cal["cameras"][cam]["T_ref_from_camera"], np.float64))
        for d in tt["detections"]:
            (x, y, z), (l, w, h), yaw = d["center"], d["size"], d["yaw"]
            c2 = _corners(x, y, yaw, l, w)[:4]
            pts = np.array([[px, py, z + dz] for dz in (-h / 2, h / 2) for px, py in c2])
            pc = pts @ Tinv[:3, :3].T + Tinv[:3, 3]
            if (pc[:, 2] < 0.5).any():
                continue
            uv = pc @ K.T
            uv = uv[:, :2] / uv[:, 2:3]
            col = tuple(int(CLS_HEX[d["label_id"]][i:i + 2], 16) for i in (1, 3, 5))
            for i, j in [(0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (5, 6), (6, 7), (7, 4), (0, 4), (1, 5), (2, 6),
                         (3, 7)]:
                p, q = np.round(uv[i]).astype(int), np.round(uv[j]).astype(int)
                cv2.line(img, tuple(p), tuple(q), (0, 0, 0), 4, cv2.LINE_AA)
                cv2.line(img, tuple(p), tuple(q), col, 2, cv2.LINE_AA)
        tiles[cam] = img
    fig = plt.figure(figsize=(12, 6.4), facecolor=SURFACE)
    gs = fig.add_gridspec(2, 3, wspace=0.04, hspace=0.12)
    pos = {"CAM_FRONT_LEFT": (0, 0), "CAM_FRONT": (0, 1), "CAM_FRONT_RIGHT": (0, 2), "CAM_BACK_LEFT": (1, 0),
           "CAM_BACK_RIGHT": (1, 2)}
    for cam, (r, c) in pos.items():
        ax = fig.add_subplot(gs[r, c])
        ax.imshow(tiles[cam])
        ax.set_axis_off()
        ax.set_title(cam, color=INK2, fontsize=8, loc="left")
    ax = fig.add_subplot(gs[1, 1])
    _bev_axes(ax, 22, "")
    ax.set_xlabel("")
    ax.set_ylabel("")
    for t in man["targets"]:
        ax.plot(t["center"][1], t["center"][0], marker="x", color=INK2, ms=7, mew=1.2, zorder=4)
    _draw_boxes(ax, rows_ref, solid=False, zorder=2)
    _draw_boxes(ax, rows_tt, solid=True)
    ax.set_ylim(-20, 18)
    ax.set_xlim(20, -20)
    _legend(ax, [Line2D([], [], color=INK2, lw=1, ls=(0, (3, 2)), label="stored CPU ref"),
                 Line2D([], [], color=INK2, marker="x", lw=0, label="pattern target")])
    ax.get_legend().remove()
    fig.legend(handles=[Line2D([], [], color=CLS_HEX[i], lw=2, label=n) for i, n in enumerate(CLASSES)]
               + [Line2D([], [], color=INK2, lw=1, ls=(0, (3, 2)), label="stored CPU reference"),
                  Line2D([], [], color=INK2, marker="x", lw=0, label="target of the pattern")],
               loc="lower center", ncol=7, fontsize=8, facecolor=SURFACE, edgecolor=GRID, labelcolor=INK2)
    fig.suptitle(f"streampetr-p150 shipped synthetic sample: {TT_LABEL}, {tt['num_detections']} boxes "
                 f"(stored CPU reference: {ref['num_detections']})", color=INK, fontsize=11, x=0.01, ha="left")
    fig.text(0.01, 0.925, "generated test pattern (code/scripts/make_sample.py), Apache-2.0; base_link, x forward",
             color=INK2, fontsize=8, ha="left")
    out = MEDIA / "streampetr_synthetic_5cam_tt.png"
    fig.savefig(out, dpi=100, facecolor=SURFACE)
    print("wrote", out, out.stat().st_size)
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    d = sub.add_parser("datasets")
    d.add_argument("--ps019")
    d.add_argument("--ns0103")
    d.add_argument("--work", default=str(ROOT / "logs" / "streampetr" / "release" / "render"))
    t = sub.add_parser("tt-vs-cpu")
    t.add_argument("--ps019", required=True)
    t.add_argument("--f1", required=True)
    s = sub.add_parser("synthetic")
    s.add_argument("--tt-body", required=True, help="a JSON with the TT /predict body (or {'body': ...})")
    a = ap.parse_args()
    return {"datasets": datasets, "tt-vs-cpu": tt_vs_cpu, "synthetic": synthetic}[a.cmd](a)


if __name__ == "__main__":
    sys.exit(main())