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 with 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'))
        # TODO: what's the difference between masking and reading .obj file
        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')
        combined_faces_int = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0])).astype(np.int32)).to('cuda')

        idx = 36000
        x = combined[idx, :].cpu().detach().numpy()
        y = np.array(info['camLoc'], dtype=np.float32)
        dist = np.sqrt(np.sum((x - y) ** 2))

        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))
        # deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined, faces=combined_faces.cpu().numpy().astype(np.int32))
        t = time.perf_counter() - start
        print(t)
        start = time.perf_counter()
        depth_image, v2i = self.render(combined, combined_faces, info['camLoc'])
        t = time.perf_counter() - start
        print(t)

        depth_image = depth_image.cpu().numpy()
        dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]]
        depth_image[v2i[idx, 1], v2i[idx, 0]] = 100
        combined = combined.cpu().numpy()

        diff = np.abs(dist - dist_img)

        import plotly.express as ex
        import plotly.graph_objects as go
        ex.imshow(deodr_depth[:,:,0]).show()
        ex.imshow(depth_image).show()

        go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
                          go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'),
                         ]).show()

        # v2i_coverage = np.zeros((480, 640), dtype=np.float32)
        # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
        # ex.imshow(v2i_coverage).show()

        # 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