import os
import numpy as np
import plotly.express as ex
import plotly.graph_objects as go
from simpleicp import PointCloud, SimpleICP
from scipy.spatial.distance import cdist
from custom_utils.custom_utils import _maskFaces



os.environ["CUDA_VISIBLE_DEVICES"] = '2'
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)

from custom_utils.smpl import SMPL


folder = '/home/workspace/data/dataset/cape_release'

def get_frames():
    # sequences = np.load()
    test = np.load(os.path.join(folder, 'sequences/03383/shortlong_ATUsquat/shortlong_ATUsquat.000105.npz'))
    faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy'))
    min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy')) 
    offsets_cano = test['v_cano'] - min_body
    pose = test['pose']
    mask = np.abs(offsets_cano.sum(axis=-1)) < 1e-2

    smpl = SMPL()
    smpl.run_garment()

    # verts, faces = _maskFaces(test['v_posed'], faces, mask)

    x = test['v_cano']
    y = min_body
    fig = go.Figure(data=[
            go.Mesh3d(
                # 8 vertices of a cube
                x=x[:,0],
                y=x[:,1],
                z=x[:,2],
                # i, j and k give the vertices of triangles
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),
            go.Mesh3d(
                # 8 vertices of a cube
                x=y[:,0],
                y=y[:,1],
                z=y[:,2],
                # i, j and k give the vertices of triangles
                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.5,1.5],),
            yaxis = dict(nticks=4, range=[-1.5,1.5],),
            zaxis = dict(nticks=4, range=[-1.5,1.5],),)
            )

    fig.show()


    print()

folder_processed = '/home/workspace/data/dataset/cape_processed/datasets/dataset_male_4clotypes'

def get_frames_processed():
    # sequences = np.load()dataset/cape_processed/datasets/dataset_male_4clotypes
    disp = np.load(os.path.join(folder_processed, 'train/train_disp.npy'))
    faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy'))
    min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy')) 


    x = verts
    fig = go.Figure(data=[
            go.Mesh3d(
                # 8 vertices of a cube
                x=x[:,0],
                y=x[:,1],
                z=x[:,2],
                # i, j and k give the vertices of triangles
                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()


    print()


folder_pop = '/home/workspace/data/dataset/cape_2/ps/scratch/ps_shared/qma/abps_data/CAPE_packed_for_POP_release/00096_shortlong/train'

def get_frames_pop():
    # sequences = np.load()dataset/cape_2/ps/scratch/ps_shared/qma/abps_data/CAPE_packed_for_POP_release/00096_shortlong/train/shortlong_ATUsquat.000009.npz
    test = np.load(os.path.join(folder_pop, 'shortlong_ATUsquat.000009.npz'))
    faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy'))
    min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy')) 
    offsets_cano = test['v_cano'] - min_body
    mask = np.abs(offsets_cano.sum(axis=-1)) < 1e-2

    # verts, faces = _maskFaces(test['v_posed'], faces, mask)

    x = test['v_cano']
    y = min_body
    fig = go.Figure(data=[
            go.Mesh3d(
                # 8 vertices of a cube
                x=x[:,0],
                y=x[:,1],
                z=x[:,2],
                # i, j and k give the vertices of triangles
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),
            go.Mesh3d(
                # 8 vertices of a cube
                x=y[:,0],
                y=y[:,1],
                z=y[:,2],
                # i, j and k give the vertices of triangles
                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()


    print()



if __name__ == '__main__':
    train_split = get_frames()

    print()