import os
import glob
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader


class HumanML3DDataset(Dataset):
    def __init__(self, data_dir='./new_joint_vecs', window_size=100):
        self.data_dir = str(data_dir)
        self.window_size = window_size
        self.file_paths = glob.glob(os.path.join(self.data_dir, '*.npy'))

        data_dir_path = Path(self.data_dir).resolve()
        base_dir = data_dir_path.parent
        self.mean = torch.from_numpy(np.load(base_dir / 'Mean.npy')).float()
        self.std = torch.from_numpy(np.load(base_dir / 'Std.npy')).float()

        print(f"Dataset initialized: found {len(self.file_paths)} files.")

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        data = np.load(self.file_paths[idx])
        tensor_data = torch.from_numpy(data).float()

        tensor_data = (tensor_data - self.mean) / self.std

        t_len, _ = tensor_data.shape
        if t_len > self.window_size:
            start = torch.randint(0, t_len - self.window_size + 1, (1,)).item()
            tensor_data = tensor_data[start : start + self.window_size, :]
        elif t_len < self.window_size:
            padding_size = self.window_size - t_len
            last_pose = tensor_data[-1].unsqueeze(0)
            padding = last_pose.repeat(padding_size, 1)
            tensor_data = torch.cat([tensor_data, padding], dim=0)

        tensor_data = tensor_data.permute(1, 0)
        return tensor_data