import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import tensorflow as tf

import sys
import shutil
import tqdm
import random
import math
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 custom_utils.custom_utils import _maskFaces

from scipy.spatial import cKDTree
from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals


# from hbertiche/PBNS
""" COLLISION FUNCTIONS START """
F = np.load("custom_utils/smpl_faces.npy")


def _nearest_neighbour(V, B):
	tree = cKDTree(B.numpy())
	return tree.query(V.numpy(), workers=-1)[1]

@tf.function
def _tf_nn_parallel(V, B):
	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])


def fix_collision(V, B, mask, thr=.004):
    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
   
    # compute correspondences
    idx = _tf_nn_parallel(V, B)
    N, M = V.shape[:2]
    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
    idx = tf.stack((_idx, idx), axis=-1)
    # compute loss
    D = V - tf.gather_nd(B, idx)
    b_vn = tf.gather_nd(B_vn, idx)
    dot = tf.einsum('abc,abc->ab', D, b_vn)
    loss = tf.minimum(dot - thr, 0)  * mask
    B_offset_closest = tf.gather_nd(B, idx) + (b_vn / 2e2)
    output = tf.where(loss[:,:,tf.newaxis] < 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()