import os
import sys
import time
import cv2
import random
import numpy as np
import tensorflow as tf

from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
from custom_utils.smpl import SMPL
import plotly.graph_objects as go
import plotly.express as ex
import cmapy as cm

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

    uv2mesh = UV2Mesh()
    smpl = SMPL()

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

    factory = DatasetFactory()
    signature_dynamic = { 'id': (),
                          'pose_shape': (72 + 11,),
                        #   'uv_posed': (3, 512, 256,),
                        #   'uv_body_posed': (3, 512, 256,),
                        #   'mesh_body_posed': (14475, 3,),
                        #   'mesh_posed': (14475, 3,),
                          'mesh_unposed': (14475, 3,),
                        #   'uv_unposed': (3, 512, 256,),
                        #   'uv_curvature': (512, 256),
                        #   'mesh_curvature':  (14475)
                        }
    signature_static = { #'uv_static': (3, 512, 256,),
                        #  'mesh_static': (14475, 3,),
                         'mesh_mask': (14475,),
                        #  'uv_mask': (512, 256,),
                        #  'uv_shaped_body': (3, 512, 256,),
                          }

    dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
                                  batch_size=1,
                                  signature=signature_dynamic,
                                  shuffle=False)
    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
                                   signature=signature_static)

    meshes = []
    current_id = ''
    count = 0
    for i, batch in enumerate(dataset):
        id = batch[0].numpy()[0].decode('utf-8')
        if id == current_id:
            continue
        current_id = id

        print(i)
        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))

        # uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
        _, f = _maskFaces(np.ones((14475, 3)), F, mask)

        # deshaped one
        ###
        pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
        gender = int(shape[0, -1] > 0)
        if gender == 0:
            continue

        shape_offsets = tf.tensordot(shape[:,:10], smpl.shapedirs[gender], axes=[[1], [2]]).numpy()
        mesh_unposed -= np.squeeze(shape_offsets)[smpl.to_keep, :]
        ###

        mesh_unposed[:,2] += (0.5 * count)
        meshes.append(go.Mesh3d(
                # 8 vertices of a cube
                x=mesh_unposed[mask,0],
                y=mesh_unposed[mask,2],
                z=mesh_unposed[mask,1],
                # i, j and k give the vertices of triangles
                i = f[:,0],
                j = f[:,1],
                k = f[:,2],
            ))

        count += 1
        if count > 10:
            break

    fig = go.Figure(data=meshes)

    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=10, z=1))
    fig.update_layout(
        scene = dict(
            xaxis = dict(nticks=4, range=[-1,1],),
                        yaxis = dict(nticks=4, range=[-1,19],),
                        zaxis = dict(nticks=4, range=[-1,1],),)
                        )

    fig.show()

    print()