import os
import sys
import tqdm
import argparse
import random
import string
import time
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 calculate_dist(folder, neighbourhood, frames):
    out = dict()
    # frames = frames[int(args.start):int(args.end)]

    for frame in tqdm.tqdm(frames):
        id = frame.split('_')[0]
        if id in out.keys():
            continue

        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
        k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
        out[id] = k_neighbors

    return out

def generate_curvatures(folder, neighbourhood, frames):

    # ids = os.listdir(folder)[int(args.start):int(args.end)]
    # frames = frames[int(args.start):int(args.end)]
    os.makedirs(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}'), exist_ok=True)

    k_neighbors_cache = calculate_dist(folder, neighbourhood, frames)

    # for frame in tqdm.tqdm(frames):

    #     id = frame.split('_')[0]
    #     frame_nr = frame.split('_')[1].split('.')[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.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
    #     _, faces = _maskFaces(np.zeros((14475,3)), F, mask) #np.load(os.path.join(folder, 'faces', f'{id}.npy'))
    #     v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][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)

    #     verts = v

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

    #     cloud = PyntCloud.from_file(temp_file)

    #     k_neighbors = k_neighbors_cache[id] # 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) + ")"])
    #     out = np.zeros((14475), dtype=np.float32)
    #     out[mask] = curvature
        
    #     np.savez_compressed(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}', f'{id}_{frame_nr}.npz'), mesh_curvature_raw=out)

    return None


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, frames, neighbourhood):

    os.makedirs(os.path.join(folder, f'mesh_curvature_{neighbourhood}_2'), exist_ok=True)
    labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)

    for i, frame in enumerate(tqdm.tqdm(frames)):
        # id = frame.split('_')[0]
        # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
        # print(i) 
        curvature = curvatures[i]

        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]

        np.savez_compressed(os.path.join(folder, f'mesh_curvature_{neighbourhood}_2', f'{frame}.npz'), mesh_curvature=frame_labels)
        # 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, frames, neighbourhood):
    curvatures_all = []
    for frame in tqdm.tqdm(frames):
        curv = np.load(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}', f'{frame}.npz'))['mesh_curvature_raw']
        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/local-cache/train_skirt')
    parser.add_argument("-s", "--start", help = "start of range", default=0)
    parser.add_argument("-e", "--end", help = "end of range", default=30000)
    parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
    args = parser.parse_args()

    is_skirt = True
    skirt_modifier = '_skirt' if is_skirt else ''
    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
    # k_neighbours(k=25)

    folder = args.folder
    frames_file = 'frames' + skirt_modifier + '_10p.txt'

    nr_of_labels = 6
    neighbourhood = 100
    temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
    if not is_skirt:
        dist = np.load('custom_data/dijkstra.npy')
    else:
        dist = np.load('custom_data/dijkstra_skirt.npy')

    print('SKIRT' if is_skirt else 'NOT SKIRT')
    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {is_skirt} -neigh {neighbourhood}')

    frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
    frames = [frame.split('.')[0] for frame in frames]

    # frames = []
    # with open(os.path.join(folder, frames_file), 'r') as f:
    #     lines = f.readlines()
    #     lines = [line.strip() for line in lines]
    # frames = lines[int(args.start):int(args.end)]

    generate_curvatures(folder, neighbourhood, frames)
    # curvatures = read_raw_curvatures(folder, frames, neighbourhood)
    # bins = get_bins(curvatures, nr_of_labels)
    # np.save(f'custom_data/train_bins_skirt_{neighbourhood}.npy', bins)
    # bins = np.load(f'custom_data/train_bins_skirt_{neighbourhood}.npy')
    # generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
    # visualize_curvature_labels(folder)

    pass