"""A tour of a real LeRobot dataset: 50 SO-101 episodes of putting a pink LEGO brick into a box.

Run (the lerobot environment, see README.md):
    ~/.virtualenvs/lerobot/bin/python classes/ml/labs/lerobot/data_tour.py

The first run downloads 86 MB from the Hugging Face Hub into ~/.cache/huggingface/lerobot.
"""

import numpy as np
import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset

REPO = "lerobot/svla_so101_pickplace"


def open_dataset(**kwargs):
    # pyav decodes the videos; torchcodec, LeRobot's default, needs a PyTorch older than the one here.
    return LeRobotDataset(REPO, video_backend="pyav", **kwargs)


def columns(dataset):
    """The whole recording as arrays: one row per frame, 30 frames a second."""
    table = dataset.hf_dataset
    return {key: np.array(table[key]) for key in ("episode_index", "frame_index", "timestamp", "observation.state", "action")}


def lag(columns, joint, max_shift=15):
    """Mean |state − action| when the state is read `shift` frames after the action, for shifts 0 … 14."""
    errors = []
    for shift in range(max_shift):
        per_episode = []
        for e in np.unique(columns["episode_index"]):
            rows = columns["episode_index"] == e
            state, action = columns["observation.state"][rows, joint], columns["action"][rows, joint]
            per_episode.append(np.abs(state[shift:] - action[:len(action) - shift]).mean())
        errors.append(float(np.mean(per_episode)))
    return errors


def normalise(value, mean, std, eps=1e-8):
    """What ACT's preprocessor does to every joint value and action before the network sees it."""
    return (value - mean) / (std + eps)


if __name__ == "__main__":
    dataset = open_dataset()
    print(f"{dataset.num_episodes} episodes, {len(dataset)} frames at {dataset.fps} fps; task: {dataset.meta.tasks.index[0]}")
    frame = dataset[0]
    for key, value in frame.items():
        print(f"  {key:26s} {tuple(value.shape) if hasattr(value, 'shape') else value}")

    data = columns(dataset)
    lengths = np.bincount(data["episode_index"])
    print(f"episode lengths: {lengths.min()} to {lengths.max()} frames, median {np.median(lengths)}")

    names = dataset.meta.features["action"]["names"]
    for j, name in enumerate(names):
        errors = lag(data, j)
        best = int(np.argmin(errors))
        print(f"  {name:18s} |state − action| {errors[0]:.2f} read at the same frame, {errors[best]:.2f} read {best} frames later")

    action = data["action"]
    for j, name in enumerate(names):
        print(f"  {name:18s} action exactly −100: {int((action[:, j] == -100).sum()):3d}, exactly 100: {int((action[:, j] == 100).sum()):3d}, "
              f"exactly 0: {int((action[:, j] == 0).sum()):3d}")

    holding = (action[:, 5] < 3) & (data["observation.state"][:, 5] > 5)      # told to close, but the jaws stay open
    print(f"gripper told to close (action < 3) but measured above 5 in {holding.sum()} frames, in "
          f"{len(np.unique(data['episode_index'][holding]))} episodes; measured there: median {np.median(data['observation.state'][holding, 5]):.1f}")

    stats = dataset.meta.stats["action"]
    value = action[0, 0]
    print(f"normalising shoulder_pan action {value:.2f}: ({value:.2f} − {stats['mean'][0]:.2f}) / {stats['std'][0]:.2f} "
          f"= {normalise(value, stats['mean'][0], stats['std'][0]):.3f}")

    # One training batch as ACT asks for it: each frame with the next 100 actions (a chunk of 3.3 s).
    chunked = open_dataset(delta_timestamps={"action": [i / dataset.fps for i in range(100)]})
    loader = torch.utils.data.DataLoader(chunked, batch_size=8, shuffle=True, generator=torch.Generator().manual_seed(0))
    batch = next(iter(loader))
    for key in ("observation.images.up", "observation.images.side", "observation.state", "action", "action_is_pad"):
        print(f"  batch {key:26s} {tuple(batch[key].shape)}")
    left = lengths[data["episode_index"]] - data["frame_index"]               # frames left in the episode, this one included
    print(f"frames whose 100-action chunk runs past the episode's end: {(left < 100).sum()} of {len(left)}")