import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = '3'

import torch
from torch.utils.data import Dataset
import numpy as np
from tqdm import tqdm
from time import time
from cloth3d.read import DataReader
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
from superSMPL.smpl_torch import SMPL
from data.pose_renderer import PoseRenderer


class Dataset(Dataset):
    def __init__(self, device='cuda'):
        self.device = device
        self.smpl = SMPL(path='superSMPL', device=device)


    def __len__(self):
        return int(1e6)
    

    def __getitem__(self, idx):
        # generate a random tenssor of size 3 with values betweeen -pi and pi
        angles = torch.rand(3) * 2 * np.pi - np.pi
        angles = angles.to(self.device)
        matrix = self.smpl.rodrigues(angles[None, None, :]).to(self.device)[0]

        data = { 
                 'angles': angles,
                 'matrix': matrix,
               }


        return data
    

if __name__ == '__main__':
    from torch.utils.data import DataLoader
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dataset = Dataset(device=device)
    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)

    for batch in train_dataloader:
        batch
        pass

    pass