"""Module for replaying recorded robot actions and movements.

This module provides functionality for loading and replaying previously recorded
robot actions and movements, allowing for demonstration and analysis of robot behavior.
"""

import time

import pyarrow as pa
from dora import Node
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset

episode = 1
dataset = LeRobotDataset("cadene/reachy2_teleop_remi")
from_index = dataset.episode_data_index["from"][episode]
to_index = dataset.episode_data_index["to"][episode]
actions = dataset.hf_dataset["action"][from_index:to_index]
states = dataset.hf_dataset["observation.state"][from_index:to_index]

images = dataset[from_index.item()]["observation.images.cam_trunk"]


node = Node()

time.sleep(1)
for state in states:
    node.send_output("agent_pos", pa.array(state.numpy()))
    time.sleep(1 / 30)