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

class DatasetSimple(tf.data.Dataset):

    def _generator(folder):
        # Opening the file
        folder = folder.decode("utf-8")
        uv_folder = os.path.join(folder, 'uv_posed') 
        pose_hape_folder = os.path.join(folder, 'poses_shapes') 
        frames = os.listdir(uv_folder)
        random.shuffle(frames)

        for frame in frames:
            id, frame_nr = frame.split('.')[0].split('_')
            uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
            uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
            pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
            uv_static = np.load(os.path.join(folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
            mesh_static = np.load(os.path.join(folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']
            mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
            mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
            uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
            mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']

            # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed

    def __new__(cls, folder, res=256):
        return tf.data.Dataset.from_generator(
            cls._generator,
            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32) ),
            args=(folder, )
        )

os.environ["CUDA_VISIBLE_DEVICES"] = "1"


if __name__ == '__main__':
    
    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)

    from losses import collision_loss as collision_penalty
    from losses import collision_penalty as collision_penalty_2
    from custom_utils.custom_utils import UV2Mesh
    import matplotlib.pyplot as plt

    def plot_3d(x, y=None):
        fig = plt.figure(figsize=(15, 15))
        ax = fig.add_subplot(111, projection='3d')
        ax.set_xlim3d(-1, 1)
        ax.set_ylim3d(-1, 1)
        ax.set_zlim3d(-1, 1)
        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.set_zlabel("Z")
        ax.grid(False)
        ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
        if y is not None:
            ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
        plt.savefig('temp.jpg')

    uv2mesh = UV2Mesh()

    for d in dataset:
        d

        v = uv2mesh(d[0])
        b = d[3]
        m = d[4]

        b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)

        c = collision_penalty(v, b, m)
        c2 = collision_penalty(v, b2, m)

        c_ = collision_penalty_2(v, b, m)
        c2_ = collision_penalty_2(v, b2, m)

        print()
        pass