import os
import tqdm
import numpy as np
import plotly.graph_objects as go


if __name__ == '__main__':
    folder = '/home/workspace/local-data/train_100/mesh_posed_body_no_orient'
    mask = np.load('custom_utils/mask_256x512.npy')
    F = np.load("custom_utils/smpl_faces.npy")

    frames = os.listdir(folder)
    visited = set()
    meshes = []
    count = 0

    for frame in tqdm.tqdm(frames):
        id, frame_nr = frame.split('_')
        frame_nr = frame_nr.split('.')[0]
        if id in visited:
            continue
        visited.add(id)
        print(id)
        mesh_posed_body_no_orient = np.load(os.path.join(folder, id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
        x = mesh_posed_body_no_orient
        # x[:, 2] += 1 * (i + 1)
        x[:, 2] += 0.5 * count
        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]
                    ))
        count += 1

        if count > 40:
            break



    fig = go.Figure(data=meshes)
    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=count/2-1, z=1))
    fig.update_layout(
        scene = dict(
            xaxis = dict(nticks=4, range=[-1,1],),
                        yaxis = dict(nticks=4, range=[-1,count/2+1],),
                        zaxis = dict(nticks=4, range=[-1,1],),)
                        )

    fig.show()

    print()