import os
import sys


os.environ["CUDA_VISIBLE_DEVICES"] = '6'
import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)


import glob
import tqdm
import copy
import numpy as np
from custom_utils.smpl import SMPL

# from custom_utils.custom_utils import UV2Mesh, _maskFaces
import plotly.graph_objects as go
# import plotly.express as ex
from scipy.spatial.transform import Rotation as R


if __name__ == '__main__':
    folder = '/home/workspace/local-data/cape_processed/train/female'
    # folder = '/home/workspace/data/dataset/test_ready'
    # uv2mesh = UV2Mesh(is_skirt=False)
    
    smpl = SMPL(use_one_hot=False, orig_smpl=True)


    F = np.load("custom_utils/smpl_faces_orig.npy")

    ids = os.listdir(os.path.join(folder, 'pose_shape'))
    ids = [id.split('.')[0] for id in ids]
    ids = set(ids)
    os.makedirs(os.path.join(folder, 'mesh_posed_no_orient'), exist_ok=True)

    for id in tqdm.tqdm(ids):

        frames = sorted(glob.glob(os.path.join(folder, 'pose_shape', id + '.*')))
        frames = ['.'.join(frame.split('/')[-1].split('.')[:-1]) + '.npz' for frame in frames]


        mesh_posed = [np.load(os.path.join(folder, 'mesh_posed',  frame)) for frame in frames]
        mesh_posed = [mesh_posed[i]['mesh_posed'] for i in range(len(frames))]
        pose_shape = [np.load(os.path.join(folder, 'pose_shape',  frame[:-1] + 'y')) for frame in frames]

        # mesh_unposed = [np.load(os.path.join(folder, 'mesh_unposed_no_shape',  frame)) for frame in frames]
        # mesh_unposed = [mesh_unposed[i]['mesh_unposed_no_shape'] for i in range(len(frames))]
        # shape = np.asarray(pose_shape)[:, 72:].astype(np.float32)
        # gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
        # shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
        # mesh_unposed_shaped = [mesh_unposed[i] + shape_offsets.numpy()[i] for i in range(shape.shape[0])]


        pose_shape_no_orient = copy.deepcopy(pose_shape)
        for i in range(len(pose_shape)):
            pose_shape_no_orient[i][:3] = 0

        rodr = smpl.run_rodrigues(np.stack([x[:72] for x in pose_shape]).astype(np.float32)).numpy()

        for i, frame in enumerate(frames):
            r = R.from_matrix(rodr[i, 0, :, :])
            # assert r.as_matrix() == rodr[frame, 0, :, :]
            out = r.apply(mesh_posed[i], inverse=True)

            # if i == 50:
            #     x = smpl.run_garment(pose_shape_no_orient[i][:72][np.newaxis,:].astype(np.float32), pose_shape_no_orient[i][72:][np.newaxis,:].astype(np.float32), mesh_unposed_shaped[i][np.newaxis,:,:].astype(np.float32))[0].numpy()
            #     y = out
            #     faces = F
            #     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 = faces[:,0],
            #                 j = faces[:,1],
            #                 k = faces[:,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()


            np.savez_compressed(os.path.join(folder, 'mesh_posed_no_orient', frame), mesh_posed_no_orient=out)