#!/usr/bin/env python3
"""Convert exported LiDAR polar NPZ frames to portable XYZ-in-metres NPZ frames."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--input", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--overwrite", action="store_true")
    a = p.parse_args()
    sources = sorted(a.input.glob("*.npz"))
    if not sources:
        raise FileNotFoundError(f"no NPZ frames in {a.input}")
    a.output.mkdir(parents=True, exist_ok=True)
    written = skipped = 0
    for index, source in enumerate(sources, 1):
        target = a.output / source.name
        if target.exists() and not a.overwrite:
            skipped += 1
            continue
        with np.load(source, allow_pickle=False) as f:
            raw = np.asarray(f["points_raw"], dtype=np.float32)
            time_ns = np.asarray(f["unix_time_ns"], dtype=np.int64)
            counter = np.asarray(f["frame_counter"], dtype=np.int32)
        distance_m = raw[:, 0] * np.float32(0.001)
        azimuth = np.deg2rad(raw[:, 1])
        altitude = np.deg2rad(raw[:, 2])
        cos_alt = np.cos(altitude)
        xyz = np.column_stack((distance_m * cos_alt * np.cos(azimuth),
                               distance_m * cos_alt * np.sin(azimuth),
                               distance_m * np.sin(altitude))).astype(np.float32, copy=False)
        np.savez_compressed(target, xyz_m=xyz, intensity=raw[:, 3].astype(np.float32, copy=False),
                            progression=raw[:, 4].astype(np.float32, copy=False),
                            unix_time_ns=time_ns, frame_counter=counter)
        written += 1
        if index % 100 == 0 or index == len(sources):
            print(f"[{index}/{len(sources)}] written={written} skipped={skipped}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())