#!/usr/bin/env python3
"""Convert a Sim2Bot raw export to RLDS, the format of Open X-Embodiment.

RLDS datasets are TensorFlow Datasets (TFDS) of episodes, each a sequence of
steps, and are what OpenVLA, Octo and the Open X-Embodiment mixtures train on.
This builds one from Sim2Bot's raw export (dataset.json, and per episode
samples.csv and one MP4 per camera) with TFDS itself, in the layout of the
Open X-Embodiment dataset template (kpertsch/rlds_dataset_builder):

  steps/
    observation/image            the first camera, JPEG (H, W, 3)
    observation/image_<camera>   each other camera
    observation/state            joints, then the gripper's opening (0 closed
                                 to 1 open) when there is one; float32
    observation/joint_velocity   the joints' measured velocities
    observation/tcp              the tool point: x y z (m), qx qy qz qw, when
                                 the robot has one
    action                       the commanded joint targets (and gripper)
    reward, discount             0 and 1 each step (demonstrations)
    is_first, is_last, is_terminal
    language_instruction         the episode's task
  episode_metadata/file_path     the episode's folder in the raw export

It writes <out>/sim2bot_dataset/<version>/, which any TFDS reader opens:

  tfds.builder_from_directory("<out>/sim2bot_dataset/1.0.0").as_dataset(split="train")

OpenVLA and Octo name a dataset's image, state and action keys in their own
configuration: image, state and action here.

Needs tensorflow, tensorflow-datasets and av (PyAV, for the MP4s):

  python -m pip install tensorflow tensorflow-datasets av
  python raw-to-rlds.py <raw export folder> <out> [--val-fraction 0.1]

Checked by loading the result with TFDS and comparing every step 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 shutil
import sys
from pathlib import Path


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


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

DATASET_VERSION = "1.0.0"


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")
    if not description.get("episodes"):
        fail("the export has no episodes")
    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:
        pictures = [frame.to_ndarray(format="rgb24") for frame in container.decode(container.streams.video[0])]
    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_keys(description: dict) -> dict[str, str]:
    """Camera key → observation key: the first camera is `image`, as the template's."""
    return {
        camera["key"]: "image" if index == 0 else f"image_{camera['key']}"
        for index, camera in enumerate(description["cameras"])
    }


def has_tcp(root: Path, description: dict) -> bool:
    header, table = read_samples(root / description["episodes"][0]["directory"] / "samples.csv")
    return bool(np.isfinite(columns(header, table, "tcp.", ["x", "y", "z", "qx", "qy", "qz", "qw"])).all())


def make_builder(root: Path, description: dict, val_fraction: float, out: Path):
    robot = description["robot"]
    joints: list[str] = robot["joints"]
    state_names = joints + (["gripper"] if robot["gripper"] else [])
    cameras = image_keys(description)
    tcp = has_tcp(root, description)
    episodes = description["episodes"]
    held = max(1, round(len(episodes) * val_fraction)) if len(episodes) > 1 and val_fraction > 0 else 0
    splits = {"train": episodes[: len(episodes) - held]}
    if held:
        splits["val"] = episodes[len(episodes) - held :]

    observation = {
        cameras[camera["key"]]: tfds.features.Image(
            shape=(camera["height"], camera["width"], 3), dtype=np.uint8, encoding_format="jpeg", doc=camera["title"]
        )
        for camera in description["cameras"]
    }
    observation["state"] = tfds.features.Tensor(
        shape=(len(state_names),), dtype=np.float32, doc=f"measured: {', '.join(state_names)}"
    )
    observation["joint_velocity"] = tfds.features.Tensor(shape=(len(joints),), dtype=np.float32, doc="measured, rad/s or m/s")
    if tcp:
        observation["tcp"] = tfds.features.Tensor(shape=(7,), dtype=np.float32, doc="x, y, z (m), qx, qy, qz, qw; world frame")

    class Sim2botDataset(tfds.core.GeneratorBasedBuilder):
        """A Sim2Bot dataset, from its raw export."""

        name = "sim2bot_dataset"
        VERSION = tfds.core.Version(DATASET_VERSION)
        RELEASE_NOTES = {DATASET_VERSION: "Converted from a Sim2Bot raw export."}

        def _info(self):
            # Not dataset_info_from_configs: it looks for README and citation
            # files beside the builder's module, which a script does not have.
            return tfds.core.DatasetInfo(
                builder=self,
                description=f"{description.get('name', 'Sim2Bot dataset')}: {robot['name']} at {description['fps']} fps",
                features=tfds.features.FeaturesDict(
                    {
                        "steps": tfds.features.Dataset(
                            {
                                "observation": tfds.features.FeaturesDict(observation),
                                "action": tfds.features.Tensor(
                                    shape=(len(state_names),), dtype=np.float32, doc=f"commanded: {', '.join(state_names)}"
                                ),
                                "reward": tfds.features.Scalar(dtype=np.float32),
                                "discount": tfds.features.Scalar(dtype=np.float32),
                                "is_first": tfds.features.Scalar(dtype=np.bool_),
                                "is_last": tfds.features.Scalar(dtype=np.bool_),
                                "is_terminal": tfds.features.Scalar(dtype=np.bool_),
                                "language_instruction": tfds.features.Text(),
                            }
                        ),
                        "episode_metadata": tfds.features.FeaturesDict({"file_path": tfds.features.Text()}),
                    }
                ),
            )

        def _split_generators(self, dl_manager):
            return {split: self._generate_examples(chosen) for split, chosen in splits.items()}

        def _generate_examples(self, chosen):
            for episode in chosen:
                directory = root / episode["directory"]
                header, table = read_samples(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)
                pose = columns(header, table, "tcp.", ["x", "y", "z", "qx", "qy", "qz", "qw"]) if tcp else None
                start = episode.get("video_start_frame", 0)
                pictures = {key: decode(directory / f"{key}.mp4", frames, start) for key in cameras}
                steps = []
                for frame in range(frames):
                    seen = {cameras[key]: images[frame] for key, images in pictures.items()}
                    seen["state"] = state[frame]
                    seen["joint_velocity"] = velocity[frame]
                    if pose is not None:
                        seen["tcp"] = pose[frame]
                    steps.append(
                        {
                            "observation": seen,
                            "action": action[frame],
                            "reward": np.float32(0),
                            "discount": np.float32(1),
                            "is_first": frame == 0,
                            "is_last": frame == frames - 1,
                            "is_terminal": frame == frames - 1,
                            "language_instruction": episode.get("task", ""),
                        }
                    )
                yield episode["directory"], {"steps": steps, "episode_metadata": {"file_path": episode["directory"]}}

    return Sim2botDataset(data_dir=str(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="a folder for the TFDS dataset")
    parser.add_argument("--val-fraction", type=float, default=0.0, help="share of episodes held out as the val split")
    args = parser.parse_args()
    root = args.raw / "raw" if (args.raw / "raw" / "dataset.json").exists() else args.raw
    description = read_description(root)
    builder = make_builder(root, description, args.val_fraction, args.out)
    builder.download_and_prepare()
    # TFDS makes a downloads folder even with nothing to download.
    downloads = args.out / "downloads"
    if downloads.is_dir() and not any(path.is_file() for path in downloads.rglob("*")):
        shutil.rmtree(downloads)
    splits = {name: split.num_examples for name, split in builder.info.splits.items()}
    print(f"wrote {args.out / builder.name / DATASET_VERSION}: {splits}")


if __name__ == "__main__":
    main()