#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Demo media of ptv3-p150: run the p150 on point clouds and save its outputs, then render bird's-eye views
(``tt_ptv3.viz``, numpy + Pillow) next to the fp32 CPU reference.

    # 1. on the device (one devrun window): the served model (serve profile seg_det, default runtime params)
    bin/devrun -t 1800 -- python code/scripts/make_demo.py run --out OUT \
        --input code/tt_ptv3/samples/test_pcd.npz --input research/ptv3/public_data/inputs/ps019_f40.pcd5.npy \
        [--min-voxels-zero research/ptv3/public_data/inputs/ns0103_k20.pcd5.npy] [--glob 'DIR/*.pcd5.npy']
    # 2. on the host: BEV renders (p150 | CPU reference | differences) and a GIF of a sequence
    python code/scripts/make_demo.py compare OUT/test_pcd.tt.npz --reference code/tt_ptv3/samples/test_pcd.reference.json \
        --points code/tt_ptv3/samples/test_pcd.npz --out media/ptv3_test_pcd_tt_vs_cpu.png --title "..."
    python code/scripts/make_demo.py gif 'OUT/ps019_f*.tt.npz' --points-dir DIR --out media/x.gif --title "..."

``run`` writes per input ``<stem>.tt.npz`` (``label_ids``, ``scores``, ``classification``, ``entropy``,
``filter_keep``: the per-point outputs in input order) and ``<stem>.tt.json`` (detections as ``/predict`` dicts,
``timing_ms``, ``meta``, the runtime params). ``.npy`` inputs with 5 columns (the research ``pcd5`` layout: x, y, z,
intensity 0-255, ring) are cut to the first 4. A reference is either a stored ``/predict`` body (``*.json``) or a
research golden (``*.golden.npz``: ``point_label15``, ``boxes`` [M, 11]).
"""
from __future__ import annotations

import argparse
import base64
import glob
import io
import json
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))

DET = ["CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN", "TRAFFIC_CONE", "BARRIER"]


def _stem(path: Path) -> str:
    name = path.name
    for suf in (".pcd5.npy", ".npz", ".npy", ".pcd"):
        if name.endswith(suf):
            return name[: -len(suf)]
    return path.stem


def _points(path: Path) -> np.ndarray:
    from tt_ptv3 import io as tio

    if path.suffix == ".npy":
        a = np.load(path, allow_pickle=False)
        return np.ascontiguousarray(a[:, :4], np.float32)
    return np.asarray(tio.load_points(str(path)).points, np.float32)


def cmd_run(a) -> int:
    from tt_ptv3 import PTv3

    jobs = [(Path(p), {}) for p in a.input] + [(Path(p), {"min_voxels": 0}) for p in a.min_voxels_zero]
    for pattern in a.glob:
        jobs += [(Path(p), {}) for p in sorted(glob.glob(pattern))]
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    kw = {"weights_dir": a.weights_dir} if a.weights_dir else {}
    t0 = time.time()
    with PTv3.from_pretrained(device_id=a.device_id, variant="seg_det", **kw) as model:
        print(f"from_pretrained {time.time() - t0:.1f} s; {model.info.get('device')}", flush=True)
        for path, params in jobs:
            o = model(_points(path), **params)
            stem = _stem(path)
            np.savez_compressed(out / f"{stem}.tt.npz", label_ids=np.asarray(o.label_ids), scores=np.asarray(o.scores),
                                classification=np.asarray(o.classification), entropy=np.asarray(o.entropy),
                                filter_keep=np.asarray(o.filter_keep))
            dets = o.to_dict().get("detections") or []
            (out / f"{stem}.tt.json").write_text(json.dumps({"input": str(path), "params": params, "detections": dets,
                                                             "timing_ms": o.timing_ms, "meta": o.meta},
                                                            default=lambda v: v.tolist() if hasattr(v, "tolist") else str(v)))
            print(f"{stem}: {len(o.label_ids)} points, {len(dets)} boxes, total {o.timing_ms.get('total', 0):.1f} ms",
                  flush=True)
    return 0


def _decode(spec):
    return np.load(io.BytesIO(base64.b64decode(spec["data"])))[spec["key"]]


def load_tt(path: Path):
    with np.load(path) as z:
        labels = z["label_ids"]
    side = json.loads(Path(str(path)[: -len(".tt.npz")] + ".tt.json").read_text())
    return labels, side["detections"]


def load_reference(path: Path):
    """``(labels per point, detection dicts)`` of a stored /predict body or a research golden."""
    if path.suffix == ".json":
        b = json.loads(path.read_text())
        return _decode(b["labels"]), b.get("detections") or []
    with np.load(path) as z:
        labels, boxes = z["point_label15"], z["boxes"]
    dets = [{"label": DET[int(r[0])], "score": float(r[1]), "center": [float(r[2]), float(r[3]), float(r[4])],
             "size": [float(r[5]), float(r[6]), float(r[7])], "yaw": float(r[8])} for r in boxes]
    return labels, dets


def cmd_compare(a) -> int:
    from tt_ptv3.viz import render_compare, save

    labels, dets = load_tt(Path(a.tt))
    ref_labels, refs = load_reference(Path(a.reference))
    xyz = _points(Path(a.points))[:, :3]
    img = render_compare(xyz, labels, ref_labels, dets, refs, title=a.title, lines=a.line, half_range=a.range,
                         px_per_m=a.px_per_m)          # None: the panels fit 960 px (2 per row)
    print(save(img, a.out), img.size, f"labels agree {float((labels == ref_labels).mean()):.4%}")
    return 0


def cmd_bev(a) -> int:
    from tt_ptv3.viz import render_bev, save

    labels, dets = load_tt(Path(a.tt))
    img = render_bev(_points(Path(a.points))[:, :3], labels, dets, title=a.title, lines=a.line, half_range=a.range,
                     px_per_m=a.px_per_m, caption="p150: per-point class + 3D boxes")
    print(save(img, a.out), img.size)
    return 0


def cmd_gif(a) -> int:
    from tt_ptv3.viz import render_bev

    frames = []
    for p in sorted(glob.glob(a.tt)):
        stem = Path(p).name[: -len(".tt.npz")]
        labels, dets = load_tt(Path(p))
        xyz = _points(Path(a.points_dir) / f"{stem}.pcd5.npy")[:, :3]
        img = render_bev(xyz, labels, dets, title=f"{a.title} ({stem})", lines=a.line, half_range=a.range,
                         px_per_m=a.px_per_m, caption="p150: per-point class + 3D boxes")
        frames.append(img.convert("P", palette=1, colors=64))      # adaptive palette: small GIF
    frames[0].save(a.out, save_all=True, append_images=frames[1:], duration=int(1000 / a.fps), loop=0, optimize=True)
    print(a.out, len(frames), "frames", Path(a.out).stat().st_size, "bytes")
    return 0


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run")
    r.add_argument("--out", required=True)
    r.add_argument("--input", action="append", default=[])
    r.add_argument("--min-voxels-zero", action="append", default=[], help="inputs below the node's 16k-voxel floor")
    r.add_argument("--glob", action="append", default=[])
    r.add_argument("--device-id", type=int, default=0)
    r.add_argument("--weights-dir")
    for name, extra in (("compare", ("--reference",)), ("bev", ())):
        c = sub.add_parser(name)
        c.add_argument("tt")
        c.add_argument("--points", required=True)
        for e in extra:
            c.add_argument(e, required=True)
        c.add_argument("--out", required=True)
        c.add_argument("--title", default="")
        c.add_argument("--line", action="append", default=[])
        c.add_argument("--range", type=float, default=40.0)
        c.add_argument("--px-per-m", type=float, default=None if name == "compare" else 8.0)
    g = sub.add_parser("gif")
    g.add_argument("tt", help="glob of <stem>.tt.npz")
    g.add_argument("--points-dir", required=True)
    g.add_argument("--out", required=True)
    g.add_argument("--title", default="")
    g.add_argument("--line", action="append", default=[])
    g.add_argument("--range", type=float, default=40.0)
    g.add_argument("--px-per-m", type=float, default=6.0)
    g.add_argument("--fps", type=float, default=5.0)
    a = ap.parse_args(argv)
    return {"run": cmd_run, "compare": cmd_compare, "bev": cmd_bev, "gif": cmd_gif}[a.cmd](a)


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