import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = '1'

import math
import time
import torch
import torchvision as tv
import numpy as np
import scipy.io as sio
import plotly.express as ex

from torch.utils.data import Dataset

from render import Render


class Dataset(Dataset):
    def __init__(self, folder):
        self.folder = folder
        self.transforms = tv.transforms.ToTensor()
        self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))

        self.render = Render(max_depth=10)
        self.body_faces = np.load('custom_data/smpl_faces.npy')
        

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


    def __getitem__(self, idx):
        path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
        mesh_posed = np.load(path)['mesh_posed']

        id = self.frames[idx].split('_')[0]
        faces = np.load(os.path.join(self.folder, 'faces', id + '.npy'))
        mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy'))
        body = np.load(os.path.join(self.folder, 'mesh_body_posed', self.frames[idx][:-1] + 'y'))

        info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat'
        info = sio.loadmat(info_folder)
        combined = np.concatenate((body, mesh_posed[mask, :]))
        combined_faces = np.concatenate((self.body_faces, faces + 14475))

        start = time.perf_counter()
        depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
        t = time.perf_counter() - start

        # image = self.transforms(image)
        return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask }


if __name__ == '__main__':
    from torch.utils.data import DataLoader
    
    device = "cuda" if torch.cuda.is_available() else "cpu"

    dataset = Dataset('/home/workspace/local-data/train')
    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True)

    for batch in train_dataloader:
        batch
        pass

    pass