#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Goldens of the port from the fp32 CPU reference (``tt_ptv3.reference``; research venv, CPU only, no device).

    OMP_NUM_THREADS=2 PYTHONPATH=code tools/research-venv/bin/python code/scripts/make_goldens.py [--frames ...]

All goldens use the **spconv-faithful** neighbour builder (13 queried taps + symmetric writes; SPEC 4.1, R11): the
earlier research golden ``research/ptv3/ref_outputs/scene10_00795_full.npz`` (26-tap emulation) is superseded.

Large (never shipped) -> ``research/ptv3/goldens/`` (``--out``):

- ``<frame>.<mode>.npz``: the research golden layout (``point_feat_0..4``, ``pred_labels``, ``pred_probs``, the 9
  ``det_*`` outputs, ``boxes`` [M, 11], ``point_label15``, ``point_prob``, ``stage_counts``) for every frame and
  both neighbour modes (``alias-first`` = the port default, PLAN.md D3; ``exact``);
- ``<frame>.alias-first.taps.npz``: per-module taps (``reference/model.py TAP_NAMES``; ``taps.`` prefix) for the
  tap frames (the shipped sample, T4 scene_10/00795, nuScenes scene-0103 k20) -- the device PCC gates' references;
- ``INDEX.json``: per golden the input sha256, the host-table digests, counts, flags, box count, file sha256.

Small -> ``<frame>.<mode>.json``: counts, flags, host-table digests, per-class voxel / point histograms and the boxes
(records) -- what the host / e2e tests check without the large files (``tests/_data.small_golden``). Only the shipped
sample's go into the bundle (``code/tt_ptv3/tests/goldens/``); the PandaSet sample's go to the git-ignored
``staging_samples_pandaset/goldens/`` (not shipped until the user approves PandaSet samples, DATASETS.md N5), the T4
and nuScenes frames' to ``research/ptv3/goldens/summaries/`` (never shipped).

Frames: the shipped sample ``test_pcd`` (Autoware's test.pcd, Apache-2.0), the PandaSet sample (PandaSet 019 f40),
the 4 T4 sample-dataset frames (SPEC 7; licence UNVERIFIED, never shipped) and nuScenes scene-0103 k20 (CC BY-NC-SA;
never shipped; below the node's 16k-voxel floor: the reference runs with the floor disabled, ``node_rejects_frame``
recorded). The 16 public frames already have research goldens
from the same reference (research/ptv3/public_data/goldens; ``test_reference_cpu.py`` checks they are reproduced).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import sys
import time
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")
T4 = WS / "assets" / "ptv3" / "t4_lidar_detection_sample" / "lidar_detection_sample_dataset"
FRAMES = {
    "test_pcd": HERE.parent / "tt_ptv3" / "samples" / "test_pcd.npz",
    "ps019_f40": WS / "research" / "ptv3" / "public_data" / "ship_sample" / "pandaset_019_f40.pcd5.npy",
    "t4_scene10_00795": T4 / "scene_10/data/LIDAR_CONCAT/00795.pcd.bin",
    "t4_scene6_00860": T4 / "scene_6/data/LIDAR_CONCAT/00860.pcd.bin",
    "t4_scene7_00270": T4 / "scene_7/data/LIDAR_CONCAT/00270.pcd.bin",
    "t4_scene7_00720": T4 / "scene_7/data/LIDAR_CONCAT/00720.pcd.bin",
    "ns0103_k20": WS / "research" / "ptv3" / "public_data" / "inputs" / "ns0103_k20.pcd5.npy",
}
TAP_FRAMES = ("ps019_f40", "t4_scene10_00795", "ns0103_k20")
# taps kept in the large goldens (the CPE / attention sub-taps and the frame-independent key position embedding are
# recomputed by the tests when needed)
SKIP_TAPS = (".cpe", ".attn", "det.key_pos_embed", "det.bev_refine", "det.stem", "det.heat", "seg.labels")
OLD_T4_GOLDEN = WS / "research" / "ptv3" / "ref_outputs" / "scene10_00795_full.npz"


def load(path: Path) -> np.ndarray:
    if path.suffix == ".npy":
        return np.load(path)
    if path.suffix == ".npz":
        from tt_ptv3 import io as tio
        return np.asarray(tio.load_points(str(path)).points, np.float32)
    return np.fromfile(path, dtype=np.float32).reshape(-1, 5)


def sha256_file(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            h.update(b)
    return h.hexdigest()


def tap_wanted(name: str) -> bool:
    return not any(s in name if s.startswith(".") else name.startswith(s) for s in SKIP_TAPS)


def hist(labels: np.ndarray, names) -> dict:
    lab = np.asarray(labels)
    h = np.bincount(lab[lab != 255].astype(np.int64), minlength=len(names))[:len(names)]
    out = {names[i]: int(h[i]) for i in range(len(names)) if h[i]}
    out["invalid"] = int((lab == 255).sum())
    return out


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("--out", default=str(WS / "research" / "ptv3" / "goldens"))
    ap.add_argument("--frames", default=",".join(FRAMES))
    ap.add_argument("--small-out", help="directory of the small goldens (default: per frame, see the docstring)")
    ap.add_argument("--modes", default="alias-first,exact")
    a = ap.parse_args(argv)

    from tt_ptv3.host.preprocess import table_digests
    from tt_ptv3.reference.config import DET_CLASSES, SEG_CLASSES
    from tt_ptv3.reference.pipeline import Reference, set_threads
    from tt_ptv3.ttaw.metrics import pcc

    set_threads(2)
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    def small_dir(fid: str) -> Path:
        if a.small_out:
            d = Path(a.small_out)
        elif fid == "test_pcd":
            d = HERE.parent / "tt_ptv3" / "tests" / "goldens"            # shipped (Apache-2.0 sample)
        elif fid.startswith("ps"):
            d = HERE.parents[1] / "staging_samples_pandaset" / "goldens"  # git-ignored until PandaSet may ship
        else:
            d = out / "summaries"                                         # workspace only (T4, nuScenes)
        d.mkdir(parents=True, exist_ok=True)
        return d

    ref = Reference(a.weights)
    index_path = out / "INDEX.json"
    index = json.loads(index_path.read_text()) if index_path.is_file() else {}
    index["reference"] = ("bundles/ptv3-p150 code/scripts/make_goldens.py (tt_ptv3.reference: fp32 torch, BN folded "
                          "as the port folds, spconv-faithful neighbour tables; torch 2 threads)")
    index.setdefault("goldens", {})
    for fid in a.frames.split(","):
        pts = load(FRAMES[fid])
        for mode in a.modes.split(","):
            t0 = time.time()
            taps: dict = {} if (fid in TAP_FRAMES and mode == "alias-first") else None
            res = ref.run(pts, neighbour_mode=mode, taps=taps, tap_filter=tap_wanted)
            g = res.golden()
            stem = f"{fid}.{mode}"
            np.savez_compressed(out / f"{stem}.npz", **g)
            if taps is not None:
                np.savez_compressed(out / f"{stem}.taps.npz", **{f"taps.{k}": v for k, v in taps.items()})
            fr = res.frame
            summary = {
                "frame": fid, "neighbour_mode": mode,
                "input": str(FRAMES[fid].relative_to(WS)) if FRAMES[fid].is_relative_to(WS) else FRAMES[fid].name,
                "input_sha256": sha256_file(FRAMES[fid]), **fr.describe(),
                "node_rejects_frame": fr.rejected,
                "seg_voxel_hist": hist(res.pred_labels, SEG_CLASSES), "seg_point_hist": hist(res.seg.labels, SEG_CLASSES),
                "boxes": [dict(r, class_name=DET_CLASSES[r["label"]]) for r in res.boxes.to_records()],
                "host_table_sha256": table_digests(fr),
                "timing_s": {k: round(v, 2) for k, v in res.timing_s.items()},
            }
            if fid == "t4_scene10_00795" and mode == "exact" and OLD_T4_GOLDEN.is_file():
                with np.load(OLD_T4_GOLDEN) as z:            # R11: the faithful tables vs the old 26-tap golden
                    summary["vs_26tap_golden"] = {
                        "point_feat_pcc": [round(pcc(res.point_feat[i], z[f"enc_out_point_feat_{i}"]), 6) for i in range(5)],
                        "seg_label_agreement": round(float(np.mean(res.pred_labels == z["seg_out_pred_labels"])), 6),
                        "spec_expected": {"point_feat_pcc_1_4": [0.99970, 0.99831, 0.99713, 0.99403], "labels": 0.9921}}
            (small_dir(fid) / f"{stem}.json").write_text(json.dumps(summary, indent=1))
            index["goldens"][stem] = {k: summary[k] for k in ("frame", "neighbour_mode", "input", "input_sha256",
                                                              "stage_counts", "node_rejects_frame")}
            index["goldens"][stem].update(boxes=len(summary["boxes"]), golden=f"{stem}.npz",
                                          golden_sha256=sha256_file(out / f"{stem}.npz"),
                                          taps=f"{stem}.taps.npz" if taps is not None else None)
            index_path.write_text(json.dumps(index, indent=1))
            print(f"{stem}: counts {fr.counts} rejects={fr.rejected} boxes {len(res.boxes)} "
                  f"taps {len(taps) if taps is not None else 0} {time.time() - t0:.1f}s"
                  + (f" vs26tap {summary['vs_26tap_golden']}" if "vs_26tap_golden" in summary else ""), flush=True)
    return 0


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