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

def get_body():
    smpl = SMPL('FeatureLines/superSMPL')
    pose = np.zeros((1, 72), dtype=np.float32)
    beta = np.zeros((1, 11), dtype=np.float32)
    R = smpl.run_rodrigues(pose)
    T, _, _ = smpl.run(R, beta, True)
    return np.squeeze(T)


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):

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

    curvatures_all = []

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

        if not os.path.isdir(garment_dir):
            continue

        # if os.path.exists(garment_dir + '/curvature_25_raw.npy'):
        #     continue

        curvatures_garment = []

        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
        # assert len (garments) == 1
        if len (garments) != 1:
            continue
        garment = garments[0]
    
        bodies = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V']
        offsetss = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V']
        mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))
        
        _, faces = _maskFaces(bodies[0, :, :], F, mask)
        v = bodies[0, mask, :] + offsetss[0, mask, :]
        m = tm.Trimesh(v, faces, process=False, validate=False)

        outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
        outliers = np.nonzero(outliers)[0]
        outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
        outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
        faces = np.concatenate((faces, outliers_faces), axis=0)

        step = 0
        for body, offsets in zip(bodies, offsetss):
            verts = body + offsets
            verts = verts[mask, :]

            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)

            k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
            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) + ")"])
            curvatures_garment.append(curvature)

            print(step)
            step += 1

        # curvatures_all.append(np.array(curvatures_garment))
        np.save(garment_dir + '/curvature_25_raw.npy', np.array(curvatures_garment))

    return curvatures_all


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, curvatures):

    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)

        if not os.path.isdir(garment_dir):
            continue

        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
        # assert len (garments) == 1
        if len (garments) != 1:
            continue
        garment = garments[0]
        mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))

        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):
    ids = os.listdir(folder)#[:10]

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

        if not os.path.isdir(garment_dir):
            continue

        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
        # assert len (garments) == 1
        if len (garments) != 1:
            continue
        garment = garments[0]
    
        frame = 10

        bodies = readPC2Frame(os.path.join(folder, id, 'Body_posed.pc16'), frame, True)
        offsetss = readPC2Frame(os.path.join(folder, id, garment + '_registered.pc16'), frame, True)
        # curvature_labels = readPC2Frame(os.path.join(folder, id, 'curvature_label_6_25.pc16'), frame, True)[:, 0]
        curvature_labels = np.load(os.path.join(folder, id, 'curvature_label_6_25.npy'), allow_pickle=True)[frame, :]
        mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))
        
        verts = offsetss
        verts, faces = _maskFaces(verts, F, mask)
        # m = tm.Trimesh(verts, faces, process=False, validate=False)

        # outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
        # outliers = np.nonzero(outliers)[0]
        # outliers_n_neighbours = get_closest_n_points(verts, outliers, 2)
        # outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
        # faces = np.concatenate((faces, outliers_faces), axis=0)

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

        # m = tm.Trimesh(verts, faces, process=False, validate=False, vertex_colors=colors_binned)
        # m.show()
        x = verts
        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,1],),
                            yaxis = dict(nticks=4, range=[-1,1],),
                            zaxis = dict(nticks=4, range=[-1,1],),)
                            )

        fig.show()

        print()


def read_raw_curvatures(folder):

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

    curvatures_all = []

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

        if not os.path.isdir(garment_dir):
            continue

        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
        # assert len (garments) == 1
        if len (garments) != 1:
            continue

        curv = np.load(garment_dir + '/curvature_25_raw.npy')
        curvatures_all.append(curv)
    
    return curvatures_all


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/test')
    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("FeatureLines/superSMPL/smpl_faces.npy")
    # k_neighbours(k=25)

    folder = args.folder
    nr_of_labels = 6
    neighbourhood = 25
    temp_file = 'temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
    dist = np.load('MeshUtils/dijsktra.npy')
    # dist = np.load('MeshUtils/dijkstra_skirt.npy')

    generate_curvatures(folder, neighbourhood)
    # curvatures = read_raw_curvatures(folder)
    # bins = get_bins(curvatures, nr_of_labels)
    # np.save('train_bins_skirt.npy', bins)
    # bins = np.load('train_bins.npy')
    # generate_curvature_labels(folder, bins, curvatures)
    # visualize_curvature_labels(folder)

    pass