#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Demo media of the card (``media/``): the p150 output, drawn over the fp32 CPU reference.

    # 1. on the device: run the p150 on inputs and keep the /predict bodies (one model load for all inputs);
    #    --golden-dir also writes <stem>.ref.json: a research golden's network outputs through this bundle's
    #    post-processing (the reference of the held-out gate, tests/test_e2e_device.py::test_heldout)
    bin/devrun -t 1800 -- python code/scripts/make_demo.py run --json-dir <dir> [--golden-dir <dir>] INPUT [INPUT ...]
    # 2. no device: the p150 boxes over the CPU reference's boxes, bird's-eye view
    python code/scripts/make_demo.py compare --input code/tt_bevfusion/samples/test_pcd.npz \\
        --tt <dir>/test_pcd.tt.json --out media/bevfusion_test_pcd_tt_vs_cpu.png
    # 3. no device: an animated bird's-eye view of a frame sequence (p150 boxes only)
    python code/scripts/make_demo.py gif --tt-dir <dir> --out media/<name>.gif INPUT [INPUT ...]
    # 4. no device: agreement of every <stem>.tt.json with its reference (JSON summary, pooled)
    python code/scripts/make_demo.py agree --json-dir <dir> [STEM ...]

``run``: ``BEVFusion.from_pretrained`` (ETH dispatch, 12x10, the traced graph; ``--weights-dir`` or the HF cache) runs
every input through the Python API, which returns exactly what ``POST /predict`` returns, and writes
``<json-dir>/<stem>.tt.json`` (the body plus the device description). Inputs are the API's point inputs (``.npz`` /
``.npy`` / ``.pcd``); an ``.npz`` whose ``points`` is a plain (M, 4) float array is read as x, y, z, time_lag (the
layout of ``research/bevfusion/public_data/inputs``).

``compare`` / ``agree``: the reference is ``--reference``, else ``<stem>.ref.json`` in the JSON directory, else
``<stem>.reference.json`` next to the input. Matching is the rule of the accuracy gates (same label, BEV centre
distance < 0.2 m, ``ttaw.metrics.match_detections``); the caption states the result. Drawing:
``tt_bevfusion.viz.render_bev`` (p150 boxes in colour, reference boxes as light outlines, unmatched boxes ringed and
tagged).

numpy + Pillow only (no device for ``compare`` / ``gif`` / ``agree``).
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np

HERE = Path(__file__).resolve().parent
CODE = HERE.parent
if str(CODE) not in sys.path:
    sys.path.insert(0, str(CODE))

FIELDS4 = ("x", "y", "z", "time_lag")
TITLE = "BEVFusion-L (Autoware autoware_bevfusion v2.0, lidar-only)"
MATCH_DIST_M = 0.2


def load_cloud(path: Path):
    """The API input of a file; an ``.npz`` / ``.npy`` holding a plain (M, 4) float array is x, y, z, time_lag."""
    from tt_bevfusion import io as tio

    if path.suffix in (".npz", ".npy"):
        if path.suffix == ".npz":
            with np.load(path, allow_pickle=False) as z:
                a = z["points"] if "points" in z.files else None
        else:
            a = np.load(path, allow_pickle=False)
        if a is not None and a.dtype.names is None and a.ndim == 2 and a.shape[1] == 4:
            return tio.PointCloud(np.ascontiguousarray(a, np.float32), FIELDS4)
    return str(path)


def xyzt(path: Path) -> np.ndarray:
    """(N, 4) x, y, z, time_lag for drawing."""
    from tt_bevfusion import io as tio

    cloud = tio.load_points(load_cloud(path))
    names = list(cloud.fields)
    xyz = cloud.select(("x", "y", "z"))
    lag = cloud.select(("time_lag",)) if "time_lag" in names else np.zeros((len(xyz), 1), np.float32)
    return np.concatenate([np.asarray(xyz, np.float32), np.asarray(lag, np.float32).reshape(-1, 1)], 1)


def rows(dets: List[Dict[str, Any]]):
    centers = np.array([d["center"][:2] for d in dets], np.float64).reshape(-1, 2)
    return centers, np.array([d["label_id"] for d in dets], np.int64), np.array([d["score"] for d in dets], np.float64)


def agreement(tt: Dict[str, Any], ref: Dict[str, Any]) -> Dict[str, Any]:
    from tt_bevfusion.ttaw.metrics import match_detections

    m = match_detections(*rows(tt["detections"]), *rows(ref["detections"]), max_dist=MATCH_DIST_M)
    out = m.to_dict()
    out.update(dev=len(tt["detections"]), ref=len(ref["detections"]),
               max_score_err=float(max(m.score_errors, default=0.0)))
    return out


def find_reference(stem: str, json_dir: Optional[Path], input_path: Optional[Path]) -> Path:
    cands = []
    if json_dir is not None:
        cands.append(json_dir / f"{stem}.ref.json")
    if input_path is not None:
        cands.append(input_path.with_name(f"{stem}.reference.json"))
    for c in cands:
        if c.is_file():
            return c
    raise SystemExit(f"no reference for {stem}: tried {[str(c) for c in cands]}")


def cmd_run(a) -> int:
    from tt_bevfusion import BEVFusion

    a.json_dir.mkdir(parents=True, exist_ok=True)
    kw = {"weights_dir": str(a.weights_dir)} if a.weights_dir else {}
    with BEVFusion.from_pretrained(device_id=a.device_id, **kw) as model:
        device = model.info.get("device")
        for path in a.inputs:
            out = model(load_cloud(path))
            body = out.to_dict()
            body["device"] = device
            (a.json_dir / f"{path.stem}.tt.json").write_text(json.dumps(body, indent=1, default=float))
            line = (f"{path.name}: {len(out)} detections, voxels {body['meta'].get('num_voxels')}, bucket "
                    f"{body['meta'].get('sparse', {}).get('bucket')}, timing_ms "
                    f"{ {k: round(v, 1) for k, v in out.timing_ms.items()} }")
            if a.golden_dir is not None:
                g = a.golden_dir / f"{path.stem}_lidar.npz"
                if g.is_file():
                    with np.load(g, allow_pickle=False) as z:
                        raw = {k: np.asarray(z[k]) for k in ("bbox_pred", "score", "label_pred")}
                    ref = model.post(raw, model.validate_params({}), model=model.MODEL_NAME).to_dict()
                    (a.json_dir / f"{path.stem}.ref.json").write_text(json.dumps(ref, indent=1, default=float))
                    ag = agreement(body, ref)
                    line += f"; vs reference {ag['matched']}/{ag['ref']} (dev {ag['dev']}), max err " \
                            f"{ag['max_score_err']:.4f}"
            print(line, flush=True)
    return 0


def cmd_compare(a) -> int:
    from tt_bevfusion.viz import render_bev

    body = json.loads(a.tt.read_text())
    ref_path = a.reference or find_reference(a.tt.name.split(".")[0], a.tt.parent, a.input)
    ref = json.loads(Path(ref_path).read_text())
    ag = agreement(body, ref)
    dev = body.get("device") or {}
    where = f"dispatch {dev.get('dispatch', 'eth')}, grid {dev.get('grid', '12x10')}"
    meta = body["meta"]
    sparse = meta.get("sparse", {})
    lines = [f"p150 ({where}) vs the fp32 CPU reference of the same network: {ag['matched']} of {ag['ref']} "
             f"reference detections matched (same label, centres < {MATCH_DIST_M} m), {ag['dev']} on the p150; "
             f"max |score difference| of matched pairs {ag['max_score_err']:.4f}.",
             f"Input: {a.describe or a.input.name}; {meta.get('densify', {}).get('num_points', 0):,} points, "
             f"{meta.get('num_voxels', 0):,} voxels, sparse bucket {sparse.get('bucket')}. Forward up, rings every "
             "10 m."]
    if a.lines:
        lines += a.lines
    img = render_bev(xyzt(a.input), body["detections"], reference=ref["detections"],
                     title=f"{TITLE} on one Blackhole p150 vs the fp32 CPU reference", lines=lines,
                     x_range=tuple(a.x_range) if a.x_range else None, y_range=tuple(a.y_range) if a.y_range else None,
                     width=a.width)
    a.out.parent.mkdir(parents=True, exist_ok=True)
    img.save(a.out, optimize=True)
    print(f"{a.out}: {json.dumps(ag)}")
    return 0


def cmd_gif(a) -> int:
    from PIL import Image

    from tt_bevfusion.viz import render_bev

    frames = []
    for i, path in enumerate(a.inputs):
        body = json.loads((a.tt_dir / f"{path.stem}.tt.json").read_text())
        label = (a.labels[i] if a.labels and i < len(a.labels) else path.stem).replace("_", " ")
        img = render_bev(xyzt(path), body["detections"], title=f"{TITLE} on one Blackhole p150: {label}",
                         lines=[f"{len(body['detections'])} detections (p150 output); {a.caption}"],
                         x_range=tuple(a.x_range), y_range=tuple(a.y_range), width=a.width)
        frames.append(img.convert("P", palette=Image.ADAPTIVE, colors=64))
    a.out.parent.mkdir(parents=True, exist_ok=True)
    frames[0].save(a.out, save_all=True, append_images=frames[1:], duration=a.ms, loop=0, optimize=True)
    print(f"{a.out}: {len(frames)} frames, {a.out.stat().st_size / 1e6:.2f} MB")
    return 0


def cmd_agree(a) -> int:
    stems = a.stems or sorted(p.name[:-len(".tt.json")] for p in a.json_dir.glob("*.tt.json"))
    per, matched, n_ref, n_dev, err = {}, 0, 0, 0, 0.0
    for stem in stems:
        body = json.loads((a.json_dir / f"{stem}.tt.json").read_text())
        ref_path = a.json_dir / f"{stem}.ref.json"
        if not ref_path.is_file():
            continue
        ag = agreement(body, json.loads(ref_path.read_text()))
        per[stem] = ag
        matched, n_ref, n_dev = matched + ag["matched"], n_ref + ag["ref"], n_dev + ag["dev"]
        err = max(err, ag["max_score_err"])
    pooled = {"frames": len(per), "matched": matched, "ref": n_ref, "dev": n_dev,
              "recall": matched / max(n_ref, 1), "precision": matched / max(n_dev, 1), "max_score_err": err}
    print(json.dumps({"pooled": pooled, "frames": per}, indent=1, default=float))
    return 0


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run", help="run the p150 on inputs (device)")
    r.add_argument("inputs", nargs="+", type=Path)
    r.add_argument("--json-dir", type=Path, required=True)
    r.add_argument("--golden-dir", type=Path, help="research public_data/golden: also write <stem>.ref.json")
    r.add_argument("--weights-dir", type=Path)
    r.add_argument("--device-id", type=int, default=0)
    c = sub.add_parser("compare", help="p150 over the CPU reference (no device)")
    c.add_argument("--input", type=Path, required=True)
    c.add_argument("--tt", type=Path, required=True, help="the p150 body (make_demo.py run)")
    c.add_argument("--reference", type=Path, help="default: <stem>.ref.json next to --tt, else <stem>.reference.json "
                                                 "next to the input")
    c.add_argument("--out", type=Path, required=True)
    c.add_argument("--describe", help="input description for the caption")
    c.add_argument("--lines", nargs="*", help="extra caption lines")
    c.add_argument("--x-range", type=float, nargs=2)
    c.add_argument("--y-range", type=float, nargs=2)
    c.add_argument("--width", type=int, default=640)
    g = sub.add_parser("gif", help="animated bird's-eye view of p150 outputs (no device)")
    g.add_argument("inputs", nargs="+", type=Path)
    g.add_argument("--tt-dir", type=Path, required=True)
    g.add_argument("--out", type=Path, required=True)
    g.add_argument("--labels", nargs="*", help="one per input; \"_\" prints as a space")
    g.add_argument("--caption", default="")
    g.add_argument("--x-range", type=float, nargs=2, default=(-40.0, 80.0))
    g.add_argument("--y-range", type=float, nargs=2, default=(-50.0, 50.0))
    g.add_argument("--width", type=int, default=560)
    g.add_argument("--ms", type=int, default=100, help="frame duration")
    s = sub.add_parser("agree", help="pooled agreement of <stem>.tt.json with <stem>.ref.json (no device)")
    s.add_argument("stems", nargs="*")
    s.add_argument("--json-dir", type=Path, required=True)
    a = ap.parse_args(argv)
    return {"run": cmd_run, "compare": cmd_compare, "gif": cmd_gif, "agree": cmd_agree}[a.cmd](a)


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