import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = '4'


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 plotly.subplots import make_subplots
from custom_utils.custom_utils import _maskFaces
from postprocessing import fix_collision

from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
from custom_utils.custom_utils import _maskFaces
from auto_encoder_curv import dice_dense

def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
    if body is not None:
        mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :]
    else:
        mesh_fixed = mesh

    if do_smoothing:
        m = tm.Trimesh(mesh_fixed, 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(mesh_fixed, outliers, 2)
        outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
        faces_fixed = np.concatenate((faces, outliers_faces), axis=0)
        m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False)
        tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4)

        mesh = m.vertices
    else:
        mesh = mesh_fixed

    return mesh

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 calculate_dist(folder, neighbourhood):
    out = dict()
    frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[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 histedges_equalN(x, nbin):
    npt = len(x)
    return np.interp(np.linspace(0, npt, nbin + 1),
                     np.arange(npt),
                     np.sort(x))


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_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
    # id = frame.split('_')[0]
    # frame_nr = int(frame.split('_')[1].split('.')[0])

    # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
    # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
    # 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 = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
    #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) + ")"])
    
    bins = np.load(f'custom_data/train_bins.npy') # get_bins(curvature, nr_of_labels)
    labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)

    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]
       
    return frame_labels


def visualize_curvature_labels(folder, results_folder, frame):
    id = frame.split('_')[0]
    frame_nr = frame.split('_')[1].split('.')[0]

    mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
    verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
    pred = np.load(os.path.join(results_folder, frame + '.npy'))
    pred = postprocess(pred, faces)

    curvature_labels_gt = generate_curvatures(verts, faces, mask, 25)
    curvature_labels_pred = generate_curvatures(pred, faces, mask, 25)
    dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
    colors_binned_all = []

    for curvature_labels in [curvature_labels_gt, curvature_labels_pred]:
        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_all.append(colors_binned)

    fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])

    fig.add_trace(go.Mesh3d(
                # 8 vertices of a cube
                x=verts[:,0],
                y=verts[:,1],
                z=verts[:,2],
                # i, j and k give the vertices of triangles
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2],
                vertexcolor=np.ones((verts.shape[0],3))*200
            ),
            row=1, col=1)

    fig.add_trace(go.Mesh3d(
                # 8 vertices of a cube
                x=verts[:,0],
                y=verts[:,1],
                z=verts[:,2],
                # i, j and k give the vertices of triangles
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2],
                vertexcolor=colors_binned_all[0]
            ),
            row=1, col=2)
    
    fig.add_trace(go.Mesh3d(
                # 8 vertices of a cube
                x=-pred[:,0],
                y=pred[:,2],
                z=pred[:,1],
                # i, j and k give the vertices of triangles
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2],
                vertexcolor=colors_binned_all[1],
                text=f'dice: {dice}'
            ),
            row=1, col=3)

    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
                    scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
                    scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
                    # scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
                    # scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
                    )
    helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
                yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
                zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
        
    fig.update_layout(scene = helper,
                        scene2 = helper,
                        scene3 = helper,
                        # scene4 = helper,
                        # scene5 = helper,
                        )


    # meshes = [go.Mesh3d(
    #             x=x[:,0],
    #             y=x[:,1],
    #             z=x[:,2],
    #             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],),)
    #                     )

    return fig, dice

import wandb

def init_wandb():
    run = wandb.init(project='plots', name='review_4', dir='/home/workspace/local-data/temp/')
    artifact = wandb.Artifact('code', type='code')
    artifact.add_file('curvature_analysis.py')
    run.log_artifact(artifact)

if __name__ == '__main__':
    init_wandb()

    is_skirt = False
    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 = '/home/workspace/local-data/train'
    results_folder = '/home/workspace/data/results/post/abl_overfit_jumpsuit_2-'

    frames_file = 'frames_jumpsuit_30k.txt'
    frames = []
    with open(os.path.join(folder, frames_file), 'r') as f:
        lines = f.readlines()
        lines = [line.strip() for line in lines]
    frames = lines

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

    i = 0
    while True:
        frame = np.random.choice(frames)
        fig, dice = visualize_curvature_labels(folder, results_folder, frame)
        wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
        i += 1
        print(frame)