import os
import sys

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

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, DepthRenderer

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from superSMPL.IO import readPC2Frame, writePC2
from utils.data_utils import Frame, read_frames_file, mask_faces
from cloth3d.read import DataReader
from cloth3d.util import quads2tris


class Dataset(Dataset):
    def __init__(self, folder, frames_file='frames_all.txt'):
        self.folder = folder
        self.transforms = tv.transforms.ToTensor()
        self.frames = read_frames_file(os.path.join(folder, frames_file))

        self.render = DepthRenderer()
        self.deodr_render : Render = Render(max_depth=10)
        self.body_faces = np.load('custom_data/smpl_faces.npy')
        self.reader = DataReader(folder, 'superSMPL')
        

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


    # TODO: make it work kwith two part outfits 
    def __getitem__(self, idx):
        frame : Frame = self.frames[idx]
        path = os.path.join(self.folder, frame.id)
        # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
        mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr)

        mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
        faces = mask_faces(self.body_faces, mask)
        # faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
        # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
        body, body_faces = self.reader.read_human(frame.id, frame.frame_nr)
        body_faces = quads2tris(body_faces)

        info = self.reader.read_info(frame.id)
        combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda')
        combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda')

        start = time.perf_counter()
        # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
        # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
        deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32))
        depth_image = self.render(combined, combined_faces, info['camLoc'])
        t = time.perf_counter() - start
        print(t)

        # 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')
    dataset = Dataset('dataset/CLOTH3D/train')
    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)

    for batch in train_dataloader:
        batch
        pass

    pass