import os
import sys

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

import math
import time
import torch
import kornia
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, get_normals_from_depth
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)

        mesh_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, mesh_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, root_joint = 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[mesh_mask, :]))).to('cuda')
        combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda')

        idx_frame = 36000
        x = combined[idx_frame, :].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))
        print(f'deodr time: {time.perf_counter() - start}')

        start = time.perf_counter()
        depth_image, depth_mask, v2i, K, camera = self.render(combined, combined_faces, info['camLoc'])
        print(f'pytorch3d time: {time.perf_counter() - start}')

        start = time.perf_counter()
        # TODO: is it sure I need this K, and not some transformed one?
        normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
        print(f'normals time: {time.perf_counter() - start}')

        root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=combined.get_device()))[0,:2]

        # Testing
        # pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True)
        # pc = pc[0,:,depth_mask].cpu().numpy()

        # 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()
        # ex.imshow(normal_image[0].cpu().numpy()).show()

        # trans = np.array(info['trans'], dtype=np.float32)[:, self.frames[idx].frame_nr]
        trans = root_joint[0,:]
        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'),
                          go.Scatter3d(x=[trans[0]], y=[trans[1]], z=[trans[2]], marker=dict(size=10, line=dict(width=0)), mode='markers'),
                         ]).show()

        # go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, 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, 
                 'mesh_mask': mesh_mask, 
                 'depth_mask': depth_mask, 
                 'normal_image': normal_image
                }


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