import os
import sys
import tqdm
import argparse
import random
import string
import numpy as np
import trimesh as tm
import cmapy as cm
import scipy as sp
import scipy.io as sio
import plotly.graph_objects as go
from pyntcloud import PyntCloud
from scipy.sparse import csr_matrix, coo_matrix
from scipy.sparse.csgraph import floyd_warshall, dijkstra

from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
from custom_utils.custom_utils import _maskFaces
# from FeatureLines.superSMPL.smpl import SMPL

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)

    except RuntimeError as e:
        print(e)
from custom_utils.smpl import SMPL

smpl = SMPL(orig_smpl=True)


def get_skirt_body():
    body = readOBJ('FeatureLines/superSMPL_skirt.obj')

    return body[0]


def export_to_ply(mesh, file):
    ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
    output_file = open(file, "wb+")
    output_file.write(ply)
    output_file.close()


def get_closest_n_points(verts, idx, n):
    dist = sp.spatial.distance.cdist(verts[idx, :], verts)
    dist[dist == 0] = 1e9
    idx_sort = np.argsort(dist, axis=-1)
    closest = idx_sort[:, :n]
    return closest


def histedges_equalN(x, nbin):
    npt = len(x)
    return np.interp(np.linspace(0, npt, nbin + 1),
                     np.arange(npt),
                     np.sort(x))



def k_neighbours(k):
    body = get_skirt_body()
    adj = np.ones((14475, 14475)) * 9999 # ???
    diagonal = np.linspace(0, 14474, 14475, dtype=np.int32)
    m = tm.Trimesh(body, F, process=False, validate=False)
    edges = m.edges_unique
    lengths = m.edges_unique_length
    adj[edges[:, 0], edges[:, 1]] = lengths
    adj[edges[:, 1], edges[:, 0]] = lengths
    adj[diagonal, diagonal] = 0

    adj = csr_matrix(adj)
    dist = dijkstra(adj, directed=False, unweighted=False, limit=0.5)

    np.save('MeshUtils/dijkstra_skirt.npy', dist)

    # dist = floydWarshall(adj, 14475)

    # idx_sort = np.argsort(dist, axis=-1)
    # closest = idx_sort[:, :k]

    pass


def generate_curvatures(folder, neighbourhood, save_folder):
    ids = os.listdir(folder)[int(args.start):int(args.end)]
    # gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True)
    faces = np.load('custom_utils/smpl_faces_orig.npy')
    k_neighbors = np.load('custom_data/dijkstra_orig_argsort_25.npy')

    os.makedirs(os.path.join(save_folder, 'curvatures_raw'), exist_ok=True)

    for id in tqdm.tqdm(ids):
        garment_dir = os.path.join(folder, id)

        sequences = os.listdir(garment_dir)

        for sequence in sequences:
            frames = os.listdir(os.path.join(garment_dir, sequence))

            for frame in frames:
                data = np.load(os.path.join(folder, id, sequence, frame))
                # pose = data['poses']
                # shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release/minimal_body_shape', id, id + '_minimal.npy'))
                # gender = [0] if gender_table[id] == 'female' else [1]
                # shape = np.concatenate((shape, gender)).astype(np.float32)
                # body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=False)[0][0]

                verts = data['v_posed'][smpl.to_keep, :]

                m = tm.Trimesh(verts, faces, process=False, validate=False)
                export_to_ply(m, temp_file)

                cloud = PyntCloud.from_file(temp_file)
                # cloud.plot(mesh=True)

                ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
                cloud.add_scalar_field("curvature", ev=ev)
                curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])

                frame_name = frame.split('.')
                assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
                frame_name = '.'.join(frame_name[:-1])

                np.save(os.path.join(save_folder, 'curvatures_raw', id + '_' + frame_name + '.npy'), np.array(curvature))
                
                print(id + '_' + frame)


def get_bins(values, nr_of_bins):
    flattened = []
    for c in values:
        flattened.extend(c.flatten())
    flattened = np.array(flattened)
    bins = histedges_equalN(flattened, nr_of_bins)
    return bins


def generate_curvature_labels(folder, bins):
    os.makedirs(os.path.join(save_folder, 'mesh_curvature'), exist_ok=True)

    frames = os.listdir(folder)#[int(args.start):int(args.end)]
    labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)


    for frame in tqdm.tqdm(frames):
        curvature = np.load(os.path.join(folder, frame))
        frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
        for j in range(bins.shape[0] - 2):
            frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
        frame_labels[curvature >= bins[-2]] = labels[-1]
    
        frame_name = frame.split('.')
        assert frame_name[-1] == 'npy', 'It\'s not an npy file!'
        frame_name = '.'.join(frame_name[:-1])
        np.savez_compressed(os.path.join(save_folder, 'mesh_curvature', frame_name + '.npz'), mesh_curvature=frame_labels)

    # ids = os.listdir(folder)#[:10]
    # labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)

    # i = 0
    # for id in tqdm.tqdm(ids):
    #     garment_dir = os.path.join(folder, id)


    #     curvature = curvatures[i]
    #     i += 1

    #     curvature_label = np.zeros((curvature.shape[0], curvature.shape[1]))
    #     frame = 0
    #     for frame in range(curvature.shape[0]):
    #         frame_labels = np.zeros((curvature.shape[1]), dtype=np.float32)
    #         for j in range(bins.shape[0] - 2):
    #             frame_labels[(curvature[frame, :] >= bins[j]) * (curvature[frame, :] < bins[j+1])] = labels[j]
    #         frame_labels[curvature[frame, :] >= bins[-2]] = labels[-1]
    #         curvature_label[frame, :] = frame_labels
    #         frame += 1

    #     unmasked_label = np.zeros((curvature.shape[0], 14475), dtype=np.float32)
    #     unmasked_label[:, mask] = curvature_label
    #     np.save(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".npy"), unmasked_label)
        # writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)


def visualize_curvature_labels(folder):
    frames = os.listdir(os.path.join(folder, 'mesh_curvature'))#[:10]
    faces = np.load('custom_utils/smpl_faces_orig.npy')

    for _ in range(10):
    
        frame = random.choice(frames)
        mesh_posed = np.load(os.path.join(folder, 'mesh_posed', frame)); mesh_posed = mesh_posed['mesh_posed']
        mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', frame)); mesh_curvature = mesh_curvature['mesh_curvature']

        colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(mesh_curvature.max()+1))])
        colors_binned = np.zeros((mesh_curvature.shape[0], 3))
        for i in range(int(mesh_curvature.max())+1):
            print(np.count_nonzero(mesh_curvature == i))
            colors_binned[mesh_curvature == i, :] = colors[i]

        x = mesh_posed
        meshes = [go.Mesh3d(
                    x=x[:,0],
                    y=x[:,2],
                    z=x[:,1],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2],
                    vertexcolor=colors_binned
                    ),]

        fig = go.Figure(data=meshes)
        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
        fig.update_layout(
            scene = dict(
                xaxis = dict(nticks=4, range=[-1.5,1.5],),
                yaxis = dict(nticks=4, range=[-1.5,1.5],),
                zaxis = dict(nticks=4, range=[-1.5,1.5],),)
                )

        fig.show()

        print()


def read_raw_curvatures(folder):

    frames = os.listdir(folder)#[int(args.start):int(args.end)]

    curvatures_all = []

    for frame in tqdm.tqdm(frames):
        curv = np.load(os.path.join(folder, frame))
        curvatures_all.append(curv)
    
    return curvatures_all


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-if", "--in_folder", help = "process the given folder", default=r'/home/workspace/local-data/cape_processed/test/male/curvatures_raw')
    parser.add_argument("-of", "--out_folder", help = "process the given folder", default=r'/home/workspace/local-data/cape_processed/test/male')
    parser.add_argument("-s", "--start", help = "start of range", default=0)
    parser.add_argument("-e", "--end", help = "end of range", default=1500)
    parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
    args = parser.parse_args()

    F = np.load("custom_utils/smpl_faces_orig.npy")
    # k_neighbours(k=25)

    folder = args.in_folder
    save_folder = args.out_folder #'/home/workspace/data/dataset/cape_processed'
    nr_of_labels = 6
    neighbourhood = 25
    temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
    # dist = np.load('custom_data/dijkstra.npy')
    # dist = np.load('MeshUtils/dijkstra_skirt.npy')

    # generate_curvatures(folder, neighbourhood, save_folder)
    # curvatures = read_raw_curvatures(folder)
    # bins = get_bins(curvatures, nr_of_labels)
    # np.save('custom_data/cape_2_train_male_bins.npy', bins)
    bins = np.load('custom_data/cape_2_train_male_bins.npy')
    generate_curvature_labels(folder, bins)
    # visualize_curvature_labels(save_folder)

    pass