import pickle
from pathlib import Path

import einops
import torch
from datasets.core import TrajectoryDataset


class YorDataset(TrajectoryDataset):
    """YOR bimanual episodes converted by utils/convert_yor.py.

    Layout under data_directory:
        actions.pth            FloatTensor [N, Tmax, A] (zero-padded)
        seq_lengths.pkl        list[int]
        obses/episode_00000.pth  uint8 Tensor [T, V, H, W, C]
        meta.json              action order / views / stride (informational)

    get_frames returns (obs, act, dummy_goal) with
        obs: FloatTensor [T, V, C, H, W] in 0..1
        act: FloatTensor [T, A]
    """

    def __init__(self, data_directory, prefetch: bool = True, views: int = 3):
        self.data_directory = Path(data_directory)
        self.actions = torch.load(self.data_directory / "actions.pth")
        with open(self.data_directory / "seq_lengths.pkl", "rb") as f:
            self.seq_lengths = pickle.load(f)
        assert len(self.actions) == len(self.seq_lengths)
        self.n = len(self.seq_lengths)
        self.views = views  # 3 = [odin, wl, wr], 1 = odin only
        for i in range(self.n):
            self.actions[i, self.seq_lengths[i]:] = 0  # redo zero padding

        self.prefetch = prefetch
        self.obses = []
        if self.prefetch:
            for i in range(self.n):
                obs = torch.load(str(self.data_directory / "obses" / f"episode_{i:05d}.pth"))
                self.obses.append(obs)  # T V H W C uint8

    def get_seq_length(self, idx):
        return self.seq_lengths[idx]

    def get_all_actions(self):
        result = []
        for i in range(self.n):
            result.append(self.actions[i, :self.seq_lengths[i]])
        return torch.cat(result, dim=0)

    def get_frames(self, idx, frames):
        if self.prefetch:
            obs = self.obses[idx][frames]
        else:
            obs = torch.load(
                str(self.data_directory / "obses" / f"episode_{idx:05d}.pth"))
            obs = obs[frames]
        obs = einops.rearrange(obs, "T V H W C -> T V C H W")[:, : self.views]
        obs = obs.float() / 255.0
        act = self.actions[idx, frames]
        dummy_goal = torch.ones([obs.shape[0], 1, 1, 1])  # T V P E
        return obs, act, dummy_goal

    def __getitem__(self, idx):
        return self.get_frames(idx, range(self.get_seq_length(idx)))

    def __len__(self):
        return self.n