import os
import sys
import time
from unicodedata import name
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
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()

    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)
    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
                                   signature=signature_static)


    uv_res = 256

    def m2uv_array(m2uv):
        result = np.zeros((14475, 2), dtype=np.int32)
        for k, v in m2uv.items():
            if len(v) == 2:
                result[k] = list(v)
            elif len(v) == 1:
                x = v.pop()
                result[k] = [x, x]
            else:
                assert False

        return result

    def unique_uv_map_mask(uv_map_discreet):
        mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
        temp = np.zeros((uv_res * 2, uv_res))
        for p in uv_map_discreet:
            temp[p[0], p[1]] += 1
        # temp = temp == 1
        for i, p in enumerate(uv_map_discreet):
            if temp[p[0], p[1]] == 1:
                mask[i] = True
            elif temp[p[0], p[1]] == 2:
                mask[i] = True
                temp[p[0], p[1]] = 0
        return mask

    def uv_to_mesh(uv, m2uv):
        res = uv.shape[0]
        mesh = np.zeros((14475, 3))
        uv_pixels = uv_map_discreet[m2uv]
        # counter = np.ones((14475))
        # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
        mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
        return mesh


    uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
    m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
    m2uv = m2uv_array(m2uv)
    uv_mask = np.load('custom_utils/mask_256x512.npy')
    curvature_label_type = '6_25'

    uv_res = 256
    uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
    uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]

    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
    uv_map_mask = unique_uv_map_mask(uv_map_discreet)
    uv_counter = np.zeros((uv_map.shape[0]))
    uv_counter[m2uv[:, 0]] = 1
    uv_counter[m2uv[:, 1]] += 1
    uv_counter = uv_counter[uv_map_mask]



    for i, batch in enumerate(dataset):
        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)


        values = mesh_unposed
        uv_values = np.zeros((uv_map.shape[0], 3))
        uv_values[m2uv[:, 0]] = values
        uv_values[m2uv[:, 1]] += values
        uv_counter = np.zeros((uv_map.shape[0]))
        uv_counter[m2uv[:, 0]] = 1
        uv_counter[m2uv[:, 1]] += 1
        uv_values /= uv_counter[:, np.newaxis]
        x = uv_values
        # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
        im = ((x - x.min()) / (x.max() - x.min())) * 255


        x = np.transpose(uv_unposed, [1,2,0])
        x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
        _, 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 = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                )]).show()


        print()