#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Build the PandaSet sample of streampetr-p150 from the research ship sample (reproducible; run once per change):

    python3 code/scripts/make_samples.py [--src research/streampetr/public_data/ship_sample/pandaset_019_f40_43]
                                         [--dest staging|package]

``--dest staging`` (default) writes into the git-ignored ``staging_samples_pandaset/`` of the bundle checkout (a
layout that mirrors ``code/tt_streampetr/``): the sample does not ship until the user approves PandaSet-derived
samples (PLAN.md 6.3, research/DATASETS.md N5). ``--dest package`` writes into ``code/tt_streampetr/`` (the swap-in,
staging_samples_pandaset/README.md). The shipped sample is the synthetic one of ``code/scripts/make_sample.py``.

Writes (relative to the destination, which mirrors ``code/tt_streampetr/``):
- ``samples/pandaset_019/<frame>_<CAMERA>.jpg``: the 20 PandaSet sequence 019 JPEGs (frames 40-43,
  the five model cameras), unmodified bytes, renamed to the model slot names; ``LICENSE.txt`` (PandaSet Dataset
  Terms, CC BY 4.0) and ``SHA256SUMS``;
- ``calib/pandaset_019.json``: the calibration preset (per camera ``P`` = ``[K | 0]`` of the
  1920x1080 image, ``T_ref_from_camera`` = camera optical frame -> base_link, ``image_size``);
- ``samples/pandaset_019_f<40..43>.json``: one request manifest per frame (images, the preset,
  per-camera stamps, ``stream`` = id, anchor CAM_FRONT stamp, ``T_world_from_ego`` = map <- base_link at that stamp;
  frame 40 resets the stream: the expected outputs are a cold start).

Source: research/streampetr/public_data (PandaSet 019 converted to the Autoware input contract; see its README for
how the extrinsics, poses and stamps were derived). Contains data from PandaSet (Scale AI and Hesai),
https://pandaset.org, licensed under CC BY 4.0 and the PandaSet Dataset Terms. Standard library + numpy only.
"""
from __future__ import annotations

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

import numpy as np

BUNDLE = Path(__file__).resolve().parents[2]
PKG = BUNDLE / "code" / "tt_streampetr"
DEFAULT_SRC = Path("/home/ubuntu/experiments/tt-models/research/streampetr/public_data/ship_sample/pandaset_019_f40_43")
SLOTS = ["CAM_FRONT", "CAM_FRONT_LEFT", "CAM_BACK_LEFT", "CAM_FRONT_RIGHT", "CAM_BACK_RIGHT"]
ATTRIBUTION = ("Contains data from PandaSet (Scale AI and Hesai), https://pandaset.org, licensed under CC BY 4.0 and the "
               "PandaSet Dataset Terms. Changes: JPEGs renamed (<frame>_<model camera slot>.jpg), unmodified bytes; "
               "calibration, poses and stamps converted to the Autoware base_link / camera-optical-frame conventions "
               "(research/streampetr/public_data). Scale AI and Hesai do not endorse this work. Cite: P. Xiao et al., "
               "PandaSet: Advanced Sensor Suite Dataset for Autonomous Driving, ITSC 2021.")


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


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", type=Path, default=DEFAULT_SRC)
    ap.add_argument("--dest", choices=("staging", "package"), default="staging")
    a = ap.parse_args()
    global PKG
    if a.dest == "staging":
        PKG = BUNDLE / "staging_samples_pandaset"
    man = json.loads((a.src / "manifest.json").read_text())
    em = np.load(a.src / "ego_motion.npz")
    slot_file = man["slots"]                                     # model slot -> manifest camera name
    out_img = PKG / "samples" / "pandaset_019"
    out_img.mkdir(parents=True, exist_ok=True)
    (PKG / "samples").mkdir(parents=True, exist_ok=True)
    sums = []
    for fr in man["frames"]:
        for s in SLOTS:
            src = a.src / fr["imgs"][slot_file[s]]
            dst = out_img / f"{fr['source_frame']}_{s}.jpg"
            shutil.copyfile(src, dst)
            sums.append(f"{sha256(dst)}  {dst.name}")
    shutil.copyfile(a.src / "LICENSE.txt", out_img / "LICENSE.txt")
    sums.append(f"{sha256(out_img / 'LICENSE.txt')}  LICENSE.txt")
    (out_img / "SHA256SUMS").write_text("\n".join(sums) + "\n")

    h, w = man["img_hw"]
    cams = {}
    for s in SLOTS:
        c = man["cams"][slot_file[s]]
        cams[s] = {"intrinsics": c["K"], "P": c["P"], "T_ref_from_camera": c["T_ego_cam"],
                   "source_camera": c["source_camera"]}
    calib = {"_doc": "Calibration preset of the shipped PandaSet sequence 019 sample (samples/pandaset_019). Per camera: "
                     "the intrinsics of the 1920x1080 image AS SENT (P = [K | 0], CameraInfo.P of a rectified "
                     "pinhole camera) and T_ref_from_camera = camera optical frame (x right, y down, z forward) -> "
                     "base_link (static extrinsics). base_link: " + man["base_link"] + ". " + ATTRIBUTION,
             "frame_id": "base_link", "image_size": [w, h], "cameras": cams}
    (PKG / "calib").mkdir(exist_ok=True)
    (PKG / "calib" / "pandaset_019.json").write_text(json.dumps(calib, indent=1) + "\n")

    for k, fr in enumerate(man["frames"]):
        f = fr["source_frame"]
        req = {"_doc": f"StreamPETR sample request: PandaSet sequence 019 frame {f} (5 cameras, 1920x1080 JPEG, "
                       "unmodified bytes). model(**tt_streampetr.load_sample(path)). 'stream' carries the frame's "
                       "anchor (CAM_FRONT) stamp and ego pose (map <- base_link at that stamp; the map frame is "
                       "PandaSet's world frame re-centred on frame 40)" + (", and resets the stream: frame 40 is a "
                       "cold start (zero temporal memory)" if k == 0 else "; send frames 40, 41, ... in order on one "
                       "stream id") + ". License: CC BY 4.0 + PandaSet Dataset Terms (pandaset_019/LICENSE.txt).",
               "images": {s: f"pandaset_019/{f}_{s}.jpg" for s in SLOTS},
               "calibration": {"preset": "pandaset_019"},
               "frame": f,
               "camera_stamps_s": {s: fr["stamps"][s] for s in SLOTS},
               "stream": {"id": "pandaset_019", "timestamp_s": fr["anchor_stamp"],
                          "T_world_from_ego": np.asarray(em["T_map_base"][k], np.float64).tolist()}}
        if k == 0:
            req["stream"]["reset"] = True
        (PKG / "samples" / f"pandaset_019_f{f}.json").write_text(json.dumps(req, indent=1) + "\n")
    print(f"wrote {len(sums) - 1} JPEGs, calib/pandaset_019.json and {len(man['frames'])} request manifests to {PKG}")


if __name__ == "__main__":
    main()