#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Build the samples of transfusion-p150 from the workspace research data (provenance in samples/README.md and
staging_samples_pandaset/README.md).

    python -I code/scripts/make_samples.py --research /home/ubuntu/experiments/tt-models/research/transfusion \
        --autoware /home/ubuntu/experiments/tt-models/upstream/autoware_universe

- ``pandaset_019_f40.npz``: ``research/transfusion/public_data/shippable/pandaset_019_f40/pandaset_019_f40_sweeps2.npy``
  (PandaSet 019 frames 40 + 39 in Autoware's densified base_link contract, float32 (M, 5)) as a structured array with
  fields x, y, z, intensity, time_lag (values unchanged; the file sha256 is checked against the sample's SHA256SUMS);
  ``PANDASET_LICENSE.txt`` is the verbatim PandaSet Dataset Terms next to it. Written to the git-ignored
  ``staging_samples_pandaset/`` of the bundle (not shipped until the user approves PandaSet samples,
  research/DATASETS.md N5);
- ``test_pcd.npz``: 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, fields x, y, z,
  intensity (128 everywhere in that file); bit-identical to the CenterPoint bundle's ``test_pcd.npz``.

numpy only; reads data, executes nothing from it.
"""
import argparse
import hashlib
import shutil
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
SAMPLES = HERE.parent / "tt_transfusion" / "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 pandaset(research: Path) -> None:
    src = research / "public_data" / "shippable" / "pandaset_019_f40"
    npy = src / "pandaset_019_f40_sweeps2.npy"
    sums = dict(reversed(line.split()) for line in (src / "SHA256SUMS").read_text().splitlines() if line.strip())
    digest = hashlib.sha256(npy.read_bytes()).hexdigest()
    if sums.get(npy.name, sums.get("./" + npy.name)) != digest:
        sys.exit(f"pandaset: {npy.name} sha256 {digest} does not match SHA256SUMS")
    pts = np.load(npy, allow_pickle=False)
    if pts.dtype != np.float32 or pts.ndim != 2 or pts.shape[1] != 5:
        sys.exit(f"pandaset: unexpected array {pts.dtype} {pts.shape}")
    PANDASET_STAGING.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(PANDASET_STAGING / "pandaset_019_f40.npz",
                        points=structured([pts[:, i] for i in range(5)], ("x", "y", "z", "intensity", "time_lag")))
    shutil.copyfile(src / "LICENSE.txt", PANDASET_STAGING / "PANDASET_LICENSE.txt")
    print(f"pandaset_019_f40.npz: {len(pts)} points, source sha256 {digest}")


def test_pcd(autoware: Path) -> None:
    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)]
    np.savez_compressed(SAMPLES / "test_pcd.npz",
                        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()}")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--research", type=Path, required=True)
    ap.add_argument("--autoware", type=Path, required=True)
    a = ap.parse_args()
    SAMPLES.mkdir(parents=True, exist_ok=True)
    pandaset(a.research)
    test_pcd(a.autoware)
    return 0


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