import os

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


import sys
import time
import cv2
import random
import numpy as np
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

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

    uv2mesh = UV2Mesh()
    smpl = SMPL(is_skirt=True, use_one_hot=True)

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

    factory = DatasetFactory()
    signature_dynamic = { 'id': (), 'frame_nr': (),
                          'pose_shape': (72 + 11,),
                        #   'uv_posed': (3, 512, 256,),
                        #   'uv_body_posed': (3, 512, 256,),
                          # 'mesh_body_posed': (14475, 3,),
                        #   'mesh_posed_body_no_orient': (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/train_skirt',
                                  batch_size=1,
                                  signature=signature_dynamic,
                                  resolutions=[512],
                                  use_no_orient=False,
                                  use_one_hot=True)
    static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
                                   signature=signature_static,
                                   resolutions=[512])


    uv_res = 256



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

        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
        pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
        mesh_unposed_shaped = batch['mesh_unposed'] + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
        mesh_reposed = smpl.run_garment(pose, shape, mesh_unposed_shaped).numpy()[0]
        mesh_posed = np.squeeze(batch['mesh_posed'].numpy())


        # x = smpl.run(smpl.run_rodrigues(pose), shape)[0][0]
        x = mesh_reposed[mask, :]
        y = mesh_posed[mask, :]
        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
        go.Figure(data=[
                go.Mesh3d(
                    x=x[:,0],
                    y=x[:,1],
                    z=x[:,2],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                ),
                go.Mesh3d(
                    x=y[:,0],
                    y=y[:,1],
                    z=y[:,2],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                )
                ]).show()

        x = mesh_unposed_shaped.numpy()[0, mask, :]
        go.Figure(data=[
                go.Mesh3d(
                    x=x[:,0],
                    y=x[:,1],
                    z=x[:,2],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                )
                ]).show()

        # pose_2 = pose.copy()
        # pose_2[:, :3] = 0
        # # x = smpl.run(smpl.run_rodrigues(pose_2), shape)[0][0]
        # x = mesh_unposed
        # y = mesh_posed_body_no_orient
        # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
        # 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 = F[:,0],
        #             j = F[:,1],
        #             k = F[:,2]
        #         ),
        #         go.Mesh3d(
        #             # 8 vertices of a cube
        #             x=y[:,0],
        #             y=y[:,1],
        #             z=y[:,2],
        #             i = F[:,0],
        #             j = F[:,1],
        #             k = F[:,2]
        #         )]).show()

        # x = mesh_posed[mask, :]
        # y = mesh_reposed[mask, :]
        # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
        # go.Figure(data=[
        #         go.Mesh3d(
        #             x=x[:,0],
        #             y=x[:,1],
        #             z=x[:,2],
        #             i = faces[:,0],
        #             j = faces[:,1],
        #             k = faces[:,2]
        #         ),
        #         go.Mesh3d(
        #             x=y[:,0],
        #             y=y[:,1],
        #             z=y[:,2],
        #             i = faces[:,0],
        #             j = faces[:,1],
        #             k = faces[:,2]
        #         )]).show()

        print()