import os
import sys
import shutil
import tqdm
import random
import math
import numpy as np
import scipy as sp
import scipy.io as sio


from custom_utils.smpl import SMPL
import plotly.graph_objects as go


def create_no_orient_body(folder):
    ids = os.listdir(folder)
    F = np.load("custom_utils/smpl_faces.npy")

    smpl = SMPL()
    for id in tqdm.tqdm(ids):
        if not os.path.isdir(os.path.join(folder, id)):
            continue

        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))

        pose, shape, gender = info['poses'], info['shape'], info['gender']

        # if gender == 1:
        #     continue 

        pose = np.transpose(pose, [1, 0]).astype(np.float32)
        pose[:, :3] = 0
        beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
        R = smpl.run_rodrigues(pose)
        T, _, _ = smpl.run(R, np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), True)
        T = T.numpy()

        # print(gender[0,0])

        # meshes = []
        # x = T[0, :, :]
        # meshes.append(go.Mesh3d(
        #             # 8 vertices of a cube
        #             x=x[:,0],
        #             y=x[:,2],
        #             z=x[:,1],
        #             # i, j and k give the vertices of triangles
        #             i = F[:,0],
        #             j = F[:,1],
        #             k = F[:,2]
        #             ))

        # fig = go.Figure(data=meshes)
        # 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()

        # frame_nr = str(0)
        # folder_2 = '/home/workspace/local-data/train_100/mesh_posed_body_no_orient'
        # mesh_posed_body_no_orient = np.load(os.path.join(folder_2, id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']

        # meshes = []
        # x = mesh_posed_body_no_orient
        # meshes.append(go.Mesh3d(
        #             # 8 vertices of a cube
        #             x=x[:,0],
        #             y=x[:,2],
        #             z=x[:,1],
        #             # i, j and k give the vertices of triangles
        #             i = F[:,0],
        #             j = F[:,1],
        #             k = F[:,2]
        #             ))

        # fig = go.Figure(data=meshes)
        # 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()

        os.makedirs(os.path.join(folder, id, 'mesh_posed_body_no_orient'), exist_ok=True)
        for i in range(pose.shape[0]):
            np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :])




if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"

    src_folder = '/home/workspace/data/dataset/train_skirt'
    # src_folder = r'D:\Work\Dataset\validation_2'
    dest_folder = '/home/workspace/data/dataset/validation' 


    create_no_orient_body('/home/workspace/data/dataset/train_100')