#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Build the STAGED PandaSet sample of frnet-p150 and its stored CPU-reference outputs (research venv: numpy, onnx,
onnxruntime; CPU only, no device). It does not ship until the user approves PandaSet-derived samples
(research/DATASETS.md N5): it goes to the git-ignored ``staging_samples_pandaset/`` of the bundle checkout (laid out
like the package: ``samples/``, ``calib/``), where the tests find it; the frozen ``sample`` gates and the PCC frame
still use it. The shipped samples are the synthetic scans of ``code/scripts/make_sample.py``.

    cd bundles/frnet-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python code/scripts/make_samples.py

Writes (from research/frnet/public_data, PandaSet 019 frame 40, CC BY 4.0 + PandaSet Dataset Terms; D18), under
``staging_samples_pandaset/``:

- ``samples/pandaset_019_40.npz``: ``points``, a structured (x, y, z, intensity) float32 array of the
  sensor-frame Pandar64 sweep, byte-identical values to ``shippable/pandaset_019_40_pandar64_xyzi.npy`` (sha256
  checked against its meta);
- ``samples/PANDASET_LICENSE.txt``: the PandaSet Dataset Terms of Use (verbatim, sha256 checked);
- ``calib/pandaset_019_40.json``: the sensor -> base_link transform at full float64 precision (the
  converted input's ``T_base_link_sensor``: its float32 cast is the research golden's crop transform bit for bit;
  the 6-decimal copy in the sample meta is not);
- ``samples/pandaset_019_40.<variant>.reference.json`` (ot128, qt128): the ``to_dict()`` of the
  reference pipeline with ONNX Runtime on the deployed ONNX (the reference of record, S:frnet:361-365), reduced to
  what the container smoke compares (``labels``, ``scores``, ``filter_keep``, npz_compressed) plus counts. The class
  ids are checked against the research golden ``public_data/golden/pandaset_019_40_<variant>.npz`` (equal).
"""
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()
CODE = HERE.parents[1]
sys.path.insert(0, str(CODE))

from tt_frnet.io import encode_array  # noqa: E402
from tt_frnet.reference.config import VARIANTS  # noqa: E402

WS = Path("/home/ubuntu/experiments/tt-models")
PUB = WS / "research" / "frnet" / "public_data"
STEM = "pandaset_019_40"
KEEP = ("model", "frame_id", "num_points", "num_points_raw", "class_names", "class_counts")


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--weights-dir", default=str(WS / "assets" / "frnet" / "hf_lidar_frnet"))
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--variants", nargs="*", default=list(VARIANTS))
    a = ap.parse_args()
    staging = CODE.parent / "staging_samples_pandaset"
    samples, calib = staging / "samples", staging / "calib"
    samples.mkdir(parents=True, exist_ok=True)
    calib.mkdir(parents=True, exist_ok=True)
    meta = json.loads((PUB / "shippable" / f"{STEM}_meta.json").read_text())
    src = PUB / "shippable" / f"{STEM}_pandar64_xyzi.npy"
    lic = PUB / "shippable" / "LICENSE.txt"
    for path in (src, lic):
        if sha256(path) != meta["files"][path.name]:
            sys.exit(f"{path}: sha256 differs from {STEM}_meta.json")
    xyzi = np.load(src, allow_pickle=False)
    pts = np.zeros(len(xyzi), dtype=[("x", "<f4"), ("y", "<f4"), ("z", "<f4"), ("intensity", "<f4")])
    for i, name in enumerate(("x", "y", "z", "intensity")):
        pts[name] = xyzi[:, i]
    np.savez_compressed(samples / f"{STEM}.npz", points=pts)
    shutil.copyfile(lic, samples / "PANDASET_LICENSE.txt")
    T = np.load(PUB / "inputs" / f"{STEM}.npz", allow_pickle=False)["T_base_link_sensor"].astype(np.float64)
    preset = {"name": STEM,
              "description": "PandaSet 019 frame 40: static Pandar64 sensor -> base_link transform (fitted mount; "
                             "base_link = the lidar-pose origin projected to the road, x forward) for the ego crop box",
              "source": "research/frnet/public_data/inputs/pandaset_019_40.npz T_base_link_sensor "
                        "(calib/pandaset_pandar64_mount.json fit)",
              "T_base_link_from_sensor": [[float(v) for v in row] for row in T]}
    (calib / f"{STEM}.json").write_text(json.dumps(preset, indent=1) + "\n")
    print(f"sample {len(pts)} points, calibration preset {STEM}", flush=True)

    from tt_frnet.reference.pipeline import FRNetReference

    for variant in a.variants:
        ref = FRNetReference(a.weights_dir, variant, backend="ort", threads=a.threads)
        run = ref.run(samples / f"{STEM}.npz",
                      calibration={"T_base_link_from_sensor": preset["T_base_link_from_sensor"]})
        gold = np.load(PUB / "golden" / f"{STEM}_{variant}.npz", allow_pickle=False)
        out = run.output
        if not (np.array_equal(run.prepared.inputs.points, gold["autoware_points"]) and
                np.array_equal(out.label_ids, gold["autoware_class_id"]) and
                np.array_equal(out.filter_keep, gold["autoware_filter_keep"])):
            sys.exit(f"{variant}: the reference output differs from the research golden")
        body = {k: v for k, v in out.to_dict().items() if k in KEEP}
        body["labels"] = encode_array(out.label_ids, fmt="npz_compressed", key="labels")
        body["scores"] = encode_array(out.scores, fmt="npz_compressed", key="scores")
        body["filter_keep"] = encode_array(out.filter_keep, fmt="npz_compressed", key="filter_keep")
        body["meta"] = {"variant": variant, "reference": "onnxruntime fp32 (deployed ONNX)",
                        "onnx_sha256": ref.weights.onnx_sha256, "num_unique_coors": run.prepared.inputs.num_voxels,
                        "cropped_by_ego_box": int((~run.prepared.inputs.keep_mask).sum())}
        path = samples / f"{STEM}.{variant}.reference.json"
        path.write_text(json.dumps(body) + "\n")
        print(f"{variant}: N {out.num_points if hasattr(out, 'num_points') else len(out)} "
              f"M {run.prepared.inputs.num_voxels} -> {path.name} ({path.stat().st_size / 1e6:.2f} MB)", flush=True)
    return 0


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