import os


import sys
import shutil
import tqdm
import random
import math
import torch
import numpy as np
import scipy as sp
import scipy.io as sio
import plotly.graph_objects as go
import plotly.express as ex
from pytorch3d.structures import Meshes

from custom_utils.custom_utils import _maskFaces
from losses.losses import _nearest_neighbour


# from hbertiche/PBNS
""" COLLISION FUNCTIONS START """
# F = np.load("smpl/smpl_faces.npy")
to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')



def fix_collision(V, B, mask, thr=.004):
    mesh = Meshes(B, F_fix_coll)
    B_vn = torch.stack(mesh.verts_normals_list())
   
    # compute correspondences
    idx = _nearest_neighbour(V, B)
    # compute loss
    B_idx = torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
    D = V - B_idx
    b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
    dot = torch.einsum('abc,abc->ab', D, b_vn)
    loss = torch.minimum(dot - thr, torch.zeros_like(dot))  * mask
    B_offset_closest = B_idx + (b_vn / 2e2)
    output = torch.where(loss[:,:,None] < 0, B_offset_closest, V)

    return output


""" COLLISION FUNCTIONS END """



if __name__ == '__main__':

    F = np.load("custom_utils/smpl_faces.npy")
    posed = np.load('/home/workspace/local-data/train/mesh_posed/00016_10.npz')['mesh_posed'][np.newaxis,:,:]
    body = np.load('/home/workspace/local-data/train/mesh_body_posed/00016_10.npy')[np.newaxis,:,:]
    mask = np.load('/home/workspace/local-data/train/mesh_mask/00016.npy')
    posed_2 = np.load('/home/workspace/local-data/train/mesh_posed/00016_15.npz')['mesh_posed'][np.newaxis,:,:]

    fixed = fix_collision(posed_2, body, mask).numpy()

    x = posed_2[0, mask,:]
    y = body[0]
    _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
    fig = go.Figure(data=[
            go.Mesh3d(
                x=x[:,0],
                y=x[:,1],
                z=x[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2],
            ),
            go.Mesh3d(
                x=y[:,0],
                y=y[:,1],
                z=y[:,2],
                i = F[:,0],
                j = F[:,1],
                k = F[:,2],
            ),
        ])

    fig.update_layout(scene_aspectmode='manual',
                    scene_aspectratio=dict(x=1, y=1, z=1))
    fig.update_layout(
        scene = dict(
            xaxis = dict(nticks=4, range=[-1,1],),
                        yaxis = dict(nticks=4, range=[-1,1],),
                        zaxis = dict(nticks=4, range=[-1,1],),)
                        )

    fig.show()

    x = fixed[0, mask,:]
    y = body[0]
    _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
    fig = go.Figure(data=[
            go.Mesh3d(
                x=x[:,0],
                y=x[:,1],
                z=x[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2],
            ),
            go.Mesh3d(
                x=y[:,0],
                y=y[:,1],
                z=y[:,2],
                i = F[:,0],
                j = F[:,1],
                k = F[:,2],
            ),
        ])

    fig.update_layout(scene_aspectmode='manual',
                    scene_aspectratio=dict(x=1, y=1, z=1))
    fig.update_layout(
        scene = dict(
            xaxis = dict(nticks=4, range=[-1,1],),
                        yaxis = dict(nticks=4, range=[-1,1],),
                        zaxis = dict(nticks=4, range=[-1,1],),)
                        )

    fig.show()

    print()