#!/usr/bin/env python3

from __future__ import annotations

import rerun as rr
from datasets import load_dataset
from PIL import Image
from tqdm import tqdm


def log_dataset_to_rerun(dataset) -> None:
    # Special time-like columns
    TIME_LIKE = {"index", "frame_id", "timestamp"}

    # Ignore these columns
    IGNORE = {"episode_data_index_from", "episode_data_index_to", "episode_id"}

    num_rows = len(dataset)
    for row_nr in tqdm(range(num_rows)):
        row = dataset[row_nr]

        # Handle time-like columns first, since they set a state (time is an index in Rerun):
        for column_name in TIME_LIKE:
            if column_name in row:
                cell = row[column_name]
                if isinstance(cell, int):
                    rr.set_time_sequence(column_name, cell)
                elif isinstance(cell, float):
                    rr.set_time_seconds(column_name, cell)  # assume seconds
                else:
                    print(f"Unknown time-like column {column_name} with value {cell}")

        # Now log actual data columns
        for column_name in dataset.column_names:
            if column_name in TIME_LIKE or column_name in IGNORE:
                continue

            cell = row[column_name]
            if isinstance(cell, Image.Image):
                rr.log(column_name, rr.Image(cell))
            elif isinstance(cell, list):
                rr.log(column_name, rr.BarChart(cell))
            elif isinstance(cell, float) or isinstance(cell, int):
                rr.log(column_name, rr.Scalar(cell))
            else:
                rr.log(column_name, rr.TextDocument(str(cell)))


def main():
    print("Loading dataset…")
    # dataset = load_dataset("lerobot/pusht", split="train")
    dataset = load_dataset("lerobot/aloha_sim_transfer_cube_human", split="train")

    print("Selecting specific episode…")
    ds_subset = dataset.filter(lambda frame: frame["episode_id"] == 3)

    print("Starting Rerun…")
    rr.init("rerun_example_lerobot", spawn=True)

    print("Logging to Rerun…")
    log_dataset_to_rerun(ds_subset)


if __name__ == "__main__":
    main()