#!/usr/bin/env python3
"""Convert a Sim2Bot raw export to HDF5 for ACT/ALOHA or robomimic.

Dataset Studio's raw export (docs/dataset-recording.md, "The raw export") is
plain files: dataset.json, and per episode samples.csv and one MP4 per camera.
This turns it into the two HDF5 layouts imitation-learning code reads, with
each camera decoded to frames, so nothing needs a video library at training
time:

  --layout act        one episode_<n>.hdf5 per episode, as ACT
                      (tonyzhaozh/act, utils.py EpisodicDataset) and the
                      ALOHA recordings lay them out:
                        /observations/qpos           (T, D) float32
                        /observations/qvel           (T, D) float32
                        /observations/images/<cam>   (T, H, W, 3) uint8
                        /action                      (T, D) float32
                        attrs: sim = True
  --layout robomimic  one dataset.hdf5, as robomimic's SequenceDataset reads:
                        /data                        attrs: total, env_args
                        /data/demo_<n>               attrs: num_samples
                          actions (N, A), rewards (N,), dones (N,)
                          obs/robot0_joint_pos, robot0_joint_vel,
                          obs/robot0_gripper_qpos (with a gripper),
                          obs/robot0_eef_pos, robot0_eef_quat (with a tool point),
                          obs/<cam>_image (N, H, W, 3) uint8
                        /mask/train, /mask/valid     demo names, for filter keys

D is the state: the arm's joints, then the gripper's opening when there is
one (0 closed to 1 open). qvel has the same width; the gripper's velocity is
not recorded and is 0. Images are gzip-compressed in the file, which h5py
undoes transparently for any reader; --no-compress stores them raw.

The original ACT computes its statistics with torch.stack over episodes, so
every episode must be the same length: --pad holds each episode's last frame
until it is as long as the longest (qvel 0 while held). Its maintained fork,
act-plus-plus (MarkFzp/act-plus-plus), takes episodes as they are.

Needs h5py, numpy and av (PyAV, for the MP4s):

  python -m pip install h5py numpy av
  python raw-to-hdf5.py <the export's raw folder> <out> --layout act

Each layout was loaded with ACT's, act-plus-plus's and robomimic's own data
loaders, every value compared with the raw export
(https://docs.sim2bot.com/user-guide/dataset-formats/).
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import sys
from pathlib import Path


def fail(message: str) -> None:
    print(f"raw-to-hdf5: {message}", file=sys.stderr)
    sys.exit(1)


try:
    import av
    import h5py
    import numpy as np
except ImportError as error:  # pragma: no cover - a message, not a test
    fail(f"{error.name} is missing: python -m pip install h5py numpy av")


def read_description(root: Path) -> dict:
    path = root / "dataset.json"
    if not path.exists():
        fail(f"no dataset.json in {root}: point at the raw export's folder (the one holding dataset.json)")
    description = json.loads(path.read_text())
    if description.get("format") != "sim2bot-raw":
        fail(f"{path} is not a Sim2Bot raw export")
    return description


def read_samples(path: Path) -> tuple[list[str], np.ndarray]:
    with path.open(newline="") as handle:
        reader = csv.reader(handle)
        header = next(reader)
        rows = [[float(value) if value != "" else math.nan for value in row] for row in reader]
    return header, np.asarray(rows, dtype=np.float64).reshape(len(rows), len(header))


def columns(header: list[str], table: np.ndarray, prefix: str, names: list[str]) -> np.ndarray:
    index = {name: at for at, name in enumerate(header)}
    missing = [name for name in names if f"{prefix}{name}" not in index]
    if missing:
        fail(f"samples.csv has no {prefix}{missing[0]} column")
    return table[:, [index[f"{prefix}{name}"] for name in names]].astype(np.float32)


def decode(path: Path, frames: int, start: int = 0) -> np.ndarray:
    """An episode's frames of an MP4 as (frames, H, W, 3) uint8: row N of samples.csv
    is the video's frame start + N. `start` is the episode's video_start_frame, 0
    unless it was trimmed in Sim2Bot's Review (its video then holds the whole
    recording)."""
    with av.open(str(path)) as container:
        stream = container.streams.video[0]
        pictures = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)]
    if len(pictures) < start + frames or (start == 0 and len(pictures) != frames):
        fail(f"{path} has {len(pictures)} frames, samples.csv {frames} from frame {start}")
    return np.stack(pictures[start : start + frames])


def image_options(compress: bool, shape: tuple[int, ...]) -> dict:
    options = {"chunks": (1, *shape[1:])}
    if compress:
        options.update(compression="gzip", compression_opts=4)
    return options


def episode_arrays(root: Path, description: dict, episode: dict) -> dict:
    robot = description["robot"]
    joints: list[str] = robot["joints"]
    state_names = joints + (["gripper"] if robot["gripper"] else [])
    header, table = read_samples(root / episode["directory"] / "samples.csv")
    frames = table.shape[0]
    state = columns(header, table, "state.", state_names)
    action = columns(header, table, "action.", state_names)
    velocity = columns(header, table, "velocity.", joints)
    qvel = np.zeros_like(state)
    qvel[:, : len(joints)] = velocity
    tcp = columns(header, table, "tcp.", ["x", "y", "z", "qx", "qy", "qz", "qw"])
    images = {
        camera["key"]: decode(root / episode["directory"] / f"{camera['key']}.mp4", frames, episode.get("video_start_frame", 0))
        for camera in description["cameras"]
    }
    return {"frames": frames, "state": state, "qvel": qvel, "velocity": velocity, "action": action, "tcp": tcp, "images": images}


def held(arrays: dict, length: int) -> dict:
    """The episode made `length` frames long by holding its last frame, still."""
    extra = length - arrays["frames"]
    if extra <= 0:
        return arrays
    repeat = lambda values: np.concatenate([values, np.repeat(values[-1:], extra, axis=0)])
    qvel = np.concatenate([arrays["qvel"], np.zeros((extra, arrays["qvel"].shape[1]), dtype=np.float32)])
    return {
        **arrays,
        "frames": length,
        "state": repeat(arrays["state"]),
        "qvel": qvel,
        "action": repeat(arrays["action"]),
        "images": {key: repeat(pictures) for key, pictures in arrays["images"].items()},
    }


def write_act(root: Path, out: Path, description: dict, compress: bool, pad: bool) -> list[Path]:
    out.mkdir(parents=True, exist_ok=True)
    written = []
    longest = max(episode["frames"] for episode in description["episodes"])
    for episode in description["episodes"]:
        arrays = episode_arrays(root, description, episode)
        if pad:
            arrays = held(arrays, longest)
        path = out / f"episode_{episode['index']}.hdf5"
        with h5py.File(path, "w") as file:
            file.attrs["sim"] = True
            file.attrs["task"] = episode.get("task", "")
            file.attrs["fps"] = description["fps"]
            # Frames recorded, before any held to pad the episode.
            file.attrs["recorded_frames"] = episode["frames"]
            observations = file.create_group("observations")
            observations.create_dataset("qpos", data=arrays["state"])
            observations.create_dataset("qvel", data=arrays["qvel"])
            group = observations.create_group("images")
            for key, pictures in arrays["images"].items():
                group.create_dataset(key, data=pictures, **image_options(compress, pictures.shape))
            file.create_dataset("action", data=arrays["action"])
        written.append(path)
    return written


def write_robomimic(root: Path, out: Path, description: dict, compress: bool, valid_fraction: float) -> list[Path]:
    out.parent.mkdir(parents=True, exist_ok=True)
    robot = description["robot"]
    joints = len(robot["joints"])
    names = []
    with h5py.File(out, "w") as file:
        data = file.create_group("data")
        total = 0
        for episode in description["episodes"]:
            arrays = episode_arrays(root, description, episode)
            name = f"demo_{episode['index']}"
            names.append(name)
            demo = data.create_group(name)
            demo.attrs["num_samples"] = arrays["frames"]
            demo.attrs["task"] = episode.get("task", "")
            demo.create_dataset("actions", data=arrays["action"])
            demo.create_dataset("rewards", data=np.zeros(arrays["frames"], dtype=np.float32))
            dones = np.zeros(arrays["frames"], dtype=np.int64)
            dones[-1] = 1
            demo.create_dataset("dones", data=dones)
            obs = demo.create_group("obs")
            obs.create_dataset("robot0_joint_pos", data=arrays["state"][:, :joints])
            obs.create_dataset("robot0_joint_vel", data=arrays["velocity"])
            if robot["gripper"]:
                obs.create_dataset("robot0_gripper_qpos", data=arrays["state"][:, joints:])
            if np.isfinite(arrays["tcp"]).all():
                obs.create_dataset("robot0_eef_pos", data=arrays["tcp"][:, :3])
                obs.create_dataset("robot0_eef_quat", data=arrays["tcp"][:, 3:])
            for key, pictures in arrays["images"].items():
                obs.create_dataset(f"{key}_image", data=pictures, **image_options(compress, pictures.shape))
            total += arrays["frames"]
        data.attrs["total"] = total
        data.attrs["env_args"] = json.dumps(
            {
                "env_name": f"Sim2Bot {robot['name']}",
                "type": 1,
                "env_kwargs": {"robot": robot["type"], "fps": description["fps"], "source": "sim2bot-raw"},
            }
        )
        # Filter keys: the last demos, a valid_fraction of them, are held out.
        held = max(1, round(len(names) * valid_fraction)) if len(names) > 1 and valid_fraction > 0 else 0
        mask = file.create_group("mask")
        mask.create_dataset("train", data=np.array(names[: len(names) - held], dtype="S"))
        mask.create_dataset("valid", data=np.array(names[len(names) - held :], dtype="S"))
    return [out]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("raw", type=Path, help="the raw export: the folder holding dataset.json")
    parser.add_argument("out", type=Path, help="act: a folder for episode_<n>.hdf5; robomimic: the .hdf5 file")
    parser.add_argument("--layout", choices=["act", "robomimic"], required=True)
    parser.add_argument("--no-compress", action="store_true", help="store images raw instead of gzip-compressed")
    parser.add_argument(
        "--pad", action="store_true", help="act: hold each episode's last frame to the longest's length, as the original ACT needs"
    )
    parser.add_argument("--valid-fraction", type=float, default=0.1, help="robomimic: share of demos held out as mask/valid")
    args = parser.parse_args()

    root = args.raw
    if (root / "raw" / "dataset.json").exists():
        root = root / "raw"
    description = read_description(root)
    if not description["episodes"]:
        fail("the export has no episodes")
    if args.layout == "act":
        written = write_act(root, args.out, description, not args.no_compress, args.pad)
    else:
        written = write_robomimic(root, args.out, description, not args.no_compress, args.valid_fraction)
    frames = sum(episode["frames"] for episode in description["episodes"])
    print(f"wrote {len(written)} file(s), {len(description['episodes'])} episode(s), {frames} frames: {args.out}")


if __name__ == "__main__":
    main()