"""RoboTwin dataset loading script for HuggingFace datasets.

Usage:
    import datasets
    ds = datasets.load_dataset("dreamzero/robotwin", split="train", streaming=True)
"""

import os
import datasets
from datasets import Video, Value, Sequence


FEATURES = datasets.Features({
    "task": Value("string"),
    "episode_index": Value("int64"),
    "frame_index": Value("int64"),
    "state": Sequence(Value("float32"), length=44),
    "action": Sequence(Sequence(Value("float32")), length=12),
    "action_mask": Sequence(Value("bool"), length=12),
    "text": Value("string"),
    "video": Video(),
})


class DreamZeroRoboTwin(datasets.GeneratorBasedBuilder):
    """RoboTwin benchmark dataset for DreamZero."""

    VERSION = datasets.Version("1.0.0")
    BUILDER_CONFIGS = [
        datasets.BuilderConfig(
            name="default",
            version=VERSION,
            description="RoboTwin dataset for DreamZero SFT",
        ),
    ]

    def _info(self):
        return datasets.DatasetInfo(
            description="RoboTwin benchmark dataset for DreamZero",
            features=FEATURES,
            homepage="https://huggingface.co/datasets/poet70/robotwin-benchmark",
            license="cc-by-4.0",
        )

    def _split_generators(self, dl_manager):
        data_dir = self.config.data_dir
        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={"data_dir": data_dir},
            ),
        ]

    def _generate_examples(self, data_dir):
        import pandas as pd

        data_dir_path = str(data_dir)
        parquet_files = sorted([
            f for f in os.listdir(data_dir_path) if f.endswith(".parquet")
        ])

        for parquet_file in parquet_files:
            df = pd.read_parquet(os.path.join(data_dir_path, parquet_file))
            for idx, row in df.iterrows():
                yield idx, {
                    "task": row["task"],
                    "episode_index": row["episode_index"],
                    "frame_index": row["frame_index"],
                    "state": row["state"],
                    "action": row["action"],
                    "action_mask": row["action_mask"],
                    "text": row["text"],
                    "video": row["video_path"],
                }