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


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_body_no_orient': (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)
    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))

        uv_posed = (np.squeeze(batch['uv_posed'].numpy()) + 1) / 2
        uv_body_posed = (np.squeeze(batch['uv_body_posed'].numpy()) + 1) / 2
        uv_unposed = (np.squeeze(batch['uv_unposed'].numpy()) + 1) / 2
        uv_curvature = np.tile(np.argmax(np.squeeze(batch['uv_curvature'].numpy()), axis=-1)[:,:,np.newaxis], [1,1,3]) / 6
        uv_static = (np.squeeze(batch['uv_static'].numpy()) + 1) / 2

        img = np.concatenate([uv_posed, uv_body_posed, uv_unposed, uv_curvature, uv_static], axis=1)

        ex.imshow(img).show()

        print()