#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Write the samples and their stored CPU-reference outputs (research venv; CPU only, no device). Provenance in
``samples/README.md`` and ``staging_samples_pandaset/README.md``.

    OMP_NUM_THREADS=2 PYTHONPATH=code tools/research-venv/bin/python code/scripts/make_samples.py \
        [--autoware upstream/autoware_universe] [--only test_pcd|pandaset]

- ``samples/test_pcd.npz`` (shipped): autoware_universe ``perception/autoware_ground_segmentation/test/data/test.pcd``
  (Apache-2.0): the non-zero returns of the organized 64 x 1024 ASCII cloud moved from ``velodyne_top`` to base_link
  by the translation (0.6, 0, 2.0) (``test_scan_ground_filter.cpp:141-153``) with Autoware's float32 sweep
  arithmetic, a structured array ``points`` with fields x, y, z, intensity (128 everywhere in that file);
  bit-identical to the CenterPoint / TransFusion / BEVFusion bundles' ``test_pcd.npz``;
- ``staging_samples_pandaset/pandaset_019_f40.npz`` (NOT shipped until the user approves PandaSet samples,
  research/DATASETS.md N5; the directory is git-ignored): ``points`` float32 (N, 4) = x, y, z, intensity 0-255,
  base_link, from ``research/ptv3/public_data/ship_sample/pandaset_019_f40.pcd5.npy`` (PandaSet 019 frame 40, CC BY
  4.0 + PandaSet Terms; the README there lists the changes) and ``PANDASET_LICENSE.txt`` (verbatim ``LICENSE.txt``);
- next to each sample ``<stem>.reference.json`` (variant ``seg_det``, serve profile ``seg_det``) and
  ``<stem>.seg.reference.json`` (variant ``seg``, serve profile ``seg``): ``PTv3Output.to_dict()`` of the fp32 CPU
  reference with the default runtime params (``neighbour_mode`` alias-first) -- what ``server/smoke_test.py``
  compares the served output with.

numpy + the bundle's CPU reference; reads data, executes nothing from it.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path

import numpy as np

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

WS = Path("/home/ubuntu/experiments/tt-models")
SHIP = WS / "research" / "ptv3" / "public_data" / "ship_sample"
SAMPLES = HERE.parent / "tt_ptv3" / "samples"
PANDASET_STAGING = HERE.parents[1] / "staging_samples_pandaset"      # git-ignored; never staged by tt-model


def structured(cols, names):
    a = np.zeros(len(cols[0]), dtype=[(n, "<f4") for n in names])
    for c, n in zip(cols, names):
        a[n] = c
    return a


def test_pcd(autoware: Path) -> Path:
    path = autoware / "perception" / "autoware_ground_segmentation" / "test" / "data" / "test.pcd"
    lines = path.read_text().split("\n")
    end = next(i for i, line in enumerate(lines) if line.startswith("DATA"))
    fields = next(line for line in lines[:end] if line.startswith("FIELDS")).split()[1:]
    a = np.array([[float(v) for v in line.split()] for line in lines[end + 1:] if line.strip()], dtype=np.float32)
    col = {f: a[:, i] for i, f in enumerate(fields)}
    xyz = np.stack([col["x"], col["y"], col["z"]], 1)
    keep = np.abs(xyz).sum(1) > 0                       # organized-cloud (0, 0, 0) fillers
    m = np.eye(4, dtype=np.float32)
    m[:3, 3] = (0.6, 0.0, 2.0)                           # base_link <- velodyne_top, identity rotation
    x, y, z = xyz[keep, 0], xyz[keep, 1], xyz[keep, 2]
    out = [((m[r, 0] * x + m[r, 1] * y) + m[r, 2] * z) + m[r, 3] for r in range(3)]
    dst = SAMPLES / "test_pcd.npz"
    np.savez_compressed(dst, points=structured(out + [col["intensity"][keep]], ("x", "y", "z", "intensity")))
    print(f"test_pcd.npz: {int(keep.sum())} of {len(xyz)} rows, sha256 of test.pcd "
          f"{hashlib.sha256(path.read_bytes()).hexdigest()}, sha256 of the npz {hashlib.sha256(dst.read_bytes()).hexdigest()}")
    return dst


def pandaset() -> Path:
    pts = np.load(SHIP / "pandaset_019_f40.pcd5.npy")[:, :4].astype(np.float32)
    PANDASET_STAGING.mkdir(parents=True, exist_ok=True)
    dst = PANDASET_STAGING / "pandaset_019_f40.npz"
    np.savez_compressed(dst, points=pts)
    shutil.copyfile(SHIP / "LICENSE.txt", PANDASET_STAGING / "PANDASET_LICENSE.txt")
    print(f"pandaset_019_f40.npz: {pts.shape}, sha256 {hashlib.sha256(dst.read_bytes()).hexdigest()}")
    return dst


def references(sample: Path, ref) -> None:
    from tt_ptv3 import io as tio
    from tt_ptv3.api import PTv3, assemble_output

    pts = np.asarray(tio.load_points(str(sample)).points, np.float32)
    res = ref.run(pts)                      # defaults: alias-first tables, full reconstruction, both heads
    params = PTv3.validate_params({})
    raw = {"pred_labels": res.pred_labels, "pred_probs": res.pred_probs, "det": res.det}
    stem = sample.with_suffix("").name
    for name, p in ((f"{stem}.reference.json", params), (f"{stem}.seg.reference.json", {**params, "run_detection": False})):
        body = assemble_output(raw, res.frame, p, model=PTv3.MODEL_NAME).to_dict()
        body["timing_ms"] = {}
        (sample.parent / name).write_text(json.dumps(body))
        print(name, "points", body["num_points"], "detections", body.get("num_detections"), body["class_counts"])


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--weights", default=str(WS / "assets" / "ptv3" / "hf_ptv3"))
    ap.add_argument("--autoware", type=Path, default=WS / "upstream" / "autoware_universe")
    ap.add_argument("--only", choices=("test_pcd", "pandaset"))
    a = ap.parse_args(argv)
    from tt_ptv3.reference.pipeline import Reference, set_threads

    set_threads(2)
    ref = Reference(a.weights)
    if a.only in (None, "test_pcd"):
        references(test_pcd(a.autoware), ref)
    if a.only in (None, "pandaset"):
        references(pandaset(), ref)
    return 0


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