import os
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
import sys
import time
import cv2
import random
import numpy as np
import tensorflow as tf
import cmapy as cm
import plotly.graph_objects as go
import plotly.express as ex

from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
from custom_utils.smpl import SMPL
from convert_to_uv_api import convert_to_uv_api
from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals


if __name__ == '__main__':

    uv2mesh = UV2Mesh()

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

    factory = DatasetFactory()
    signature_dynamic = { 'id': (), 'frame_nr': (),
                        #   'pose_shape': (72 + 11,),
                          'uv_posed': (512, 256, 3),
                          'uv_body_posed': (512, 256, 3),
                          'mesh_body_posed': (14475, 3,),
                          'mesh_posed': (14475, 3,),
                          'mesh_unposed': (14475, 3,),
                          'uv_unposed': (512, 256, 3),
                          'uv_curvature': (512, 256, 7),
                        #   'mesh_curvature':  (14475)
                        }
    signature_static = { 'uv_static': (512, 256, 3),
                         '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/train',
                                  batch_size=1,
                                  signature=signature_dynamic,
                                  testing=True,
                                  shuffle=True)
    static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
                                   signature=signature_static)


    uv_res = 256



    for i, batch in enumerate(dataset):
        print(i)
        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))

        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)

        # vertex_normals(batch['mesh_posed'], tf.tile(F[None], [1, 1, 1]))
        # vertex_normals(batch['mesh_body_posed'], tf.tile(F[None], [1, 1, 1]))
        # vertex_normals(batch['mesh_unposed'], tf.tile(F[None], [1, 1, 1]))
        # vertex_normals(batch['mesh_static'], tf.tile(F[None], [1, 1, 1]))


        uv_norm_posed = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_posed'], tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4
        uv_norm_body_posed = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_body_posed'], tf.tile(F[None], [1, 1, 1])).numpy()), np.ones((14475), bool)) + 2) / 4
        uv_norm_unposed = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_unposed'], tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4
        uv_curvature = np.tile(np.argmax(np.squeeze(batch['uv_curvature'].numpy()), axis=-1)[:,:,np.newaxis], [1,1,3]) / 6
        uv_norm_static = (convert_to_uv_api(np.squeeze(vertex_normals(batch['mesh_static'], tf.tile(F[None], [1, 1, 1])).numpy()) * mask[:,np.newaxis], mask) + 2) / 4

        img = np.concatenate([uv_norm_posed, uv_norm_body_posed, uv_norm_unposed, uv_curvature, uv_norm_static], axis=1)

        ex.imshow(img).show()

        print()