import os
import sys

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

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

from torch.utils.data import Dataset


sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.render import Render, DepthRenderer
from superSMPL.IO import readPC2Frame, writePC2
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
from cloth3d.read import DataReader
from cloth3d.util import quads2tris, zRotMatrix
import kornia
import plotly.graph_objects as go
from data.render import intrinsic
from simpleicp import SimpleICP, PointCloud


def depth_to_3d(depth, depth_mask):
    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
    pc_full = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
    pc = [pc_full[i,:,depth_mask[i]] for i in range(len(depth_mask))]
    # x = torch.sum(torch.stack(depth_mask), dim=0).cpu().numpy() == len(depth_mask)
    # choice = [0,0]
    # while x[choice[0], choice[1]] == False:
    #     choice = (np.random.random(2) * [480, 640]).astype(int)
    # for i in range(len(pc)):
    #     pc[i] = pc[i] - pc_full[i][:,choice[0], choice[1]][:,None]
    return pc


def perspective_test(folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, fix_camera=True):
    transforms = tv.transforms.ToTensor()
    frames = read_frames_file(os.path.join(folder, frames_file))
    render = DepthRenderer()
    body_faces = np.load('custom_data/smpl_faces.npy')
    reader = DataReader(folder, 'superSMPL')

    frame : Frame = frames[0]

    path = os.path.join(folder, frame.id)
    info = reader.read_info(frame.id)
    # info['zrot'] = torch.pi / 2 + torch.pi/4

    zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
    trans = info['trans'][:,frame.frame_nr].astype(np.float32)
    # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)

    mesh_posed = reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot)
    mesh_posed = torch.tensor(mesh_posed, device=device)
    # mesh_posed_no_trans = reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, None, zrot)
    # mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=device)

    # raw_mask = sio.loadmat(os.path.join(path, frame.garment + '_mask.mat'))['mask'].reshape((-1))
    # mesh_mask = raw_mask > 0 #np.load(os.path.join(path, frame.garment + '_mask.npy'))
    # TODO: what's the difference between masking and reading .obj file. is the obj file the not registered version?
    # faces = mask_faces(self.body_faces, mesh_mask)
    faces = quads2tris(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, joints, joints_no_trans = reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=True)
    body_no_trans, _, _, _, _ = reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
    body = torch.tensor(body.astype(np.float32), device=device)
    body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=device)
    body_faces = quads2tris(body_faces)

    joints = torch.tensor(joints[0], device=device).to(torch.float32)

    offsets = [[1, -0.9, 0], [1, 0.9, 0], [1, 0, 0.4], [1, 0, -0.4], [1,-0.9,0.4], [1, 0.9, -0.4]]
    # x: - = closer, + = farther (+-0.5)
    # y: - = to the right, + = to the left (+-1)
    # z: - = up, + = down (+-0.5)

    combineds = []
    for offset in offsets:
        combineds.append(torch.concat((body, mesh_posed)) - torch.tensor(offset, device=body.get_device())[None,:])
    # combined_no_trans = torch.concat((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))
    combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(device)

    depth_images = []
    depth_masks = []
    for i, combined in enumerate(combineds):
        start = time.perf_counter()
        cam = info['camLoc'].copy()
        if not fix_camera:
            cam = cam - offsets[i]
        depth_image, depth_mask, _, _, _, _ = render(combined, combined_faces, cam)
        print(f'pytorch3d time: {time.perf_counter() - start}')
        depth_images.append(depth_image)
        depth_masks.append(depth_mask)

    # depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
    # depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask

    # rgb_image = tv.io.read_image(os.path.join(folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32).to(device)

    all_depth = torch.stack(depth_images)
    pointclouds = depth_to_3d(all_depth[:,None,:,:], depth_masks)

    all_depth_merged = torch.sum(all_depth, dim=0)
    ex.imshow(all_depth_merged.cpu()).show()

    mesh_data = []
    for pc in pointclouds:
        pc = pc.cpu().numpy()
        mesh_data.append(go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),)

    go.Figure(data = mesh_data).show()

    pc_0 = pointclouds[0] - torch.mean(pointclouds[0], dim=1)[:,None]
    pc_fix = PointCloud(torch.permute(pc_0,[1,0]).cpu().numpy(), columns=["x", "y", "z"])
    pc_registered = [np.transpose(pc_0.cpu().numpy(), (1,0))]
    for pc in pointclouds[1:]:
        pc = pc - torch.mean(pc, dim=1)[:,None]
        # pc_reg = torch.permute(pc,[1,0]).cpu().numpy()
        pc_mov = PointCloud(torch.permute(pc,[1,0]).cpu().numpy(), columns=["x", "y", "z"])
        icp = SimpleICP()
        icp.add_point_clouds(pc_fix, pc_mov)
        _, pc_reg, _, _ = icp.run(max_iterations=50, correspondences=10000, min_change=0.5, neighbors=50)
        pc_registered.append(pc_reg)
        # pass

    mesh_data = []
    for pc in pc_registered:
        pc = np.transpose(pc, (1,0))
        mesh_data.append(go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),)

    go.Figure(data = mesh_data).show()
    pass


if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"
    fix_camera = True
    perspective_test('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', fix_camera=fix_camera, device=device)


    pass