import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'

import torch
from torch.nn import Module
import numpy as np
from superSMPL.smpl_torch import SMPL
import trimesh
import time
import plotly.express as ex
import plotly.graph_objects as go


from data_utils import intrinsic, extrinsic, get_camera

from pytorch3d.renderer import (
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    SoftPhongShader,
    TexturesVertex,
    HardPhongShader,
    AmbientLights
)
from pytorch3d.structures import Meshes


class PoseRenderer(Module):
    def __init__(self, device='cuda') -> None:
        super().__init__()
        colors = np.load('custom_data/pose_colors.npy')
        self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
        self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()


    def forward(self, verts: torch.Tensor, faces: torch.Tensor):
        cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512)

        raster_settings = RasterizationSettings(
                            image_size=(512, 512), 
                            blur_radius=0.0, 
                            faces_per_pixel=1, 
                            bin_size=0
                        )
        
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=HardPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device))
        )

        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))

        return output
    

def get_dijsktra(m: trimesh.Trimesh) -> np.ndarray:
    edge_lengths = m.edges_unique_length
    # create and adjacency matrix based on the edge_lengths
    adjacency_matrix = np.zeros((m.vertices.shape[0], m.vertices.shape[0]), dtype=np.float32)
    adjacency_matrix[m.edges_unique[:,0], m.edges_unique[:,1]] = edge_lengths
    adjacency_matrix[m.edges_unique[:,1], m.edges_unique[:,0]] = edge_lengths
    # calculate the distance between all vertices using dijsktra algorithm
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import dijkstra
    dists = csr_matrix(adjacency_matrix)
    dists = dijkstra(dists, directed=False)    

    return dists


def get_distances(vertices, edges):
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import dijkstra
    num_vertices = vertices.shape[0]

    # edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:,1]], axis=-1)
    edge_lengths = np.ones(edges.shape[0])
    adjacency_matrix = csr_matrix((edge_lengths.flatten(), (edges[:, 0], edges[:, 1])), shape=(num_vertices, num_vertices))

    # Compute the distances between all vertices using Dijkstra's algorithm
    # distances = dijkstra(adjacency_matrix, directed=False)
    distances = dijkstra(adjacency_matrix, directed=False, unweighted=True)

    return distances


def get_distances2(vertices, faces):
    from scipy.sparse import coo_matrix, csgraph

    num_vertices = vertices.shape[0]
    num_faces = faces.shape[0]

    # Calculate the edge lengths
    edge_lengths = np.linalg.norm(vertices[faces[:, np.newaxis]] - vertices[faces], axis=2)

    # Create the adjacency matrix as a COO matrix
    row_indices = np.repeat(faces.flatten(), 3)
    col_indices = np.tile(faces.flatten(), 3)
    data = np.tile(edge_lengths.flatten(), 3)
    adjacency_matrix = coo_matrix((data, (row_indices, col_indices)), shape=(num_vertices, num_vertices))

    # Compute the distances between all vertices using Dijkstra's algorithm
    distances = csgraph.dijkstra(adjacency_matrix, directed=False)

    return distances


import networkx as nx
from scipy.spatial.distance import cdist
 
# def order_mesh_vertices(vertices, distances, edges):
#     # Create a graph with vertices as nodes and distances as edge weights
#     graph = nx.Graph()
#     graph.add_nodes_from(range(len(vertices)))
#     # calcualate the edge lengths
#     edge_lengths = np.array([np.linalg.norm(vertices[edges[i, 0]] - vertices[edges[i, 1]]) for i in range(edges.shape[0])])

#     # Add edges to the graph based on the mesh edges
#     for edge, edge_length in zip(edges, edge_lengths):
#         graph.add_edge(edge[0], edge[1], weight=edge_length)
 
#     ordered_vertices = []
#     ordered_indices = []
 
#     start_vertex = 0
#     # Perform breadth-first search starting from the selected vertex
#     visited = set()
#     queue = [(1221, 1221)]  # Tuple of (vertex, index)

#     while queue:
#         vertex, index = queue.pop(0)
#         if vertex not in visited:
#             visited.add(vertex)
#             ordered_vertices.append(vertices[vertex])
#             ordered_indices.append(vertex)
#             # neighbors = sorted(graph.neighbors(vertex), key=lambda n: graph.edges[vertex, n]['weight'])
#             neighbors = graph.neighbors(vertex)
#             queue.extend([(n, n) for n in neighbors])
 
#     return np.array(ordered_vertices), np.array(ordered_indices)

from collections import deque
def order_mesh_vertices(vertices, distances, edges):
    # Create a graph with vertices as nodes and distances as edge weights
    graph = nx.Graph()
    graph.add_nodes_from(range(len(vertices)))
    
    # Calculate the edge lengths
    edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:, 1]], axis=1)
    
    # Add edges to the graph based on the mesh edges
    for edge, edge_length in zip(edges, edge_lengths):
        graph.add_edge(edge[0], edge[1], weight=edge_length)
 
    ordered_indices = []
 
    start_vertex = 13833 
    # Perform breadth-first search starting from the selected vertex
    visited = set()
    queue = deque([start_vertex])

    while queue:
        vertex = queue.popleft()
        if vertex not in visited:
            visited.add(vertex)
            ordered_indices.append(vertex)
            neighbors = sorted(graph.neighbors(vertex), key=lambda n: graph.edges[vertex, n]['weight'])
            # neighbors = graph.neighbors(vertex)
            queue.extend(neighbors)
 
    return np.array(ordered_indices)


def interpolate_colors(distance_matrix, colored_indices, colored_colors):
    from scipy.spatial import distance, cKDTree
    all_colors = np.array(np.meshgrid(np.arange(0,256), np.arange(0,256), np.arange(0,256))).T.reshape(-1,3)
    color_map = dict()
    for idx, color in enumerate(all_colors):
        color_map[tuple(color)] = color
    color_map_array = np.array(list(color_map.values()))

    num_vertices = distance_matrix.shape[0]
    kd_tree = cKDTree(color_map_array)

    # Create an array to store the interpolated colors
    interpolated_colors = np.zeros((num_vertices, 3), dtype=np.uint8)

    # Assign the given colors to the corresponding vertices
    for index, color in zip(colored_indices, colored_colors):
        interpolated_colors[index] = color

    used_colors = set([tuple(color) for color in colored_colors])

    # Interpolate colors for the remaining vertices
    for vertex in range(num_vertices):
        print(vertex)
        if vertex not in colored_indices:
            # start = time.time()
            # Find the distances from the current vertex to the colored vertices
            distances = distance_matrix[vertex, colored_indices]
            # take the three closest one only
            closest_indices = np.argsort(distances)[:3]
            distances = distances[closest_indices]
            # print(time.time() - start)

            # Calculate the weights for color interpolation based on inverse distances
            weights = 1.0 / distances
            # Normalize the weights to sum up to 1
            weights /= np.sum(weights)
            # Interpolate the colors based on the weights
            color = np.dot(weights, np.array(colored_colors)[closest_indices]).astype(np.uint8)

            # get the closest color to this from color_map if this color is not available
            if tuple(color) in used_colors:
                start = time.time()
                color_idx = kd_tree.query(color, workers=48)[1]
                color = color_map_array[color_idx]
                # if color is already in used_colors then get the next closest color until find one that is not in used_colors
                k = 2
                while tuple(color) in used_colors:
                    neighbour_dist, color_idx = kd_tree.query(color, k=k, workers=48)
                    color = color_map_array[color_idx[-1]]
                    k += 1
                # assert neighbour_dist[-1] < 5
                # color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
                print(time.time() - start)

            # start = time.time()
            interpolated_colors[vertex] = color
            used_colors.add(tuple(color))
            # color_map.pop(tuple(color))
            # color_map_array = np.delete(color_map_array, np.where(np.all(color_map_array == color, axis=1))[0], axis=0)
            # np.delete(color_map_array, np.where , axis=0)
            # print(time.time() - start)

    return interpolated_colors


def interpolate_colors_2(distance_matrix, colored_indices, colored_colors):
    from scipy.spatial import distance, cKDTree
    all_colors = np.array(np.meshgrid(np.arange(0,256), np.arange(0,256), np.arange(0,256))).T.reshape(-1,3)
    color_map = dict()
    for idx, color in enumerate(all_colors):
        color_map[tuple(color)] = color
    color_map_array = np.array(list(color_map.values()))

    num_vertices = distance_matrix.shape[0]
    kd_tree = cKDTree(color_map_array)

    # Create an array to store the interpolated colors
    interpolated_colors = np.zeros((num_vertices, 3), dtype=np.uint8)

    # Assign the given colors to the corresponding vertices
    for index, color in zip(colored_indices, colored_colors):
        interpolated_colors[index] = color

    # Interpolate colors for the remaining vertices
    for vertex in range(num_vertices):
        print(vertex)
        if vertex not in colored_indices:
            # Find the distances from the current vertex to the colored vertices
            distances = distance_matrix[vertex, colored_indices]

            # Calculate the weights for color interpolation based on inverse distances
            weights = 1.0 / distances

            # Normalize the weights to sum up to 1
            weights /= np.sum(weights)

            # Interpolate the colors based on the weights
            color = np.dot(weights, colored_colors).astype(np.uint8)

            # get the closest color to this from color_map if this color is not available
            if tuple(color) not in color_map:
                kd_tree = cKDTree(color_map_array)
                color_idx = kd_tree.query(color, workers=48)[1]
                color = color_map_array[color_idx]
                # color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]

            interpolated_colors[vertex] = color
            color_map.pop(tuple(color))
            color_map_array = np.delete(color_map_array, np.where(np.all(color_map_array == color, axis=1))[0], axis=0)
            # np.delete(color_map_array, np.where , axis=0)

    return interpolated_colors


def get_edges(faces):
    edges = []
    unique_edges = set()

    # Iterate over each face
    for face in faces:
        # Get the vertices of the current face
        v0, v1, v2 = face

        # Define the three edges of the current face
        # and add them to the unique_edges set
        edge0 = tuple(sorted((v0, v1)))
        edge1 = tuple(sorted((v1, v2)))
        edge2 = tuple(sorted((v2, v0)))

        unique_edges.add(edge0)
        unique_edges.add(edge1)
        unique_edges.add(edge2)

    # Convert the unique_edges set to a list
    edges = list(unique_edges)

    # Convert the list of edges to a numpy array
    edges = np.array(edges)

    return edges


def get_data():
    from plyfile import PlyData
    # Read the .ply file
    plydata = PlyData.read('custom_data/body_color_init.ply')

    # Get the vertex coordinates
    vertices = plydata['vertex']
    coordinates = np.vstack((vertices['x'], vertices['y'], vertices['z'])).T

    # Get the vertex colors
    colors = np.vstack((vertices['red'], vertices['green'], vertices['blue'])).T

    return coordinates, colors



def find_corresponding_vertices(mesh1_vertices, mesh2_vertices, indices_mesh1):
    corresponding_vertices_mesh2 = []
    # mesh1_vertices[:, [1, 2]] = mesh1_vertices[:, [2, 1]]
    for index_mesh1 in indices_mesh1:
        vertex_mesh1 = mesh1_vertices[index_mesh1]
        for i, vertex_mesh2 in enumerate(mesh2_vertices):
            if np.allclose(vertex_mesh1, vertex_mesh2):
                corresponding_vertices_mesh2.append(i)
                break
    return np.array(corresponding_vertices_mesh2)

def transform_selected_indices(selected_indices, new_order):
    mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(new_order)}
    transformed_indices = np.array([mapping[idx] for idx in selected_indices])
    return transformed_indices

def transform_faces(faces, new_order):
    mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(new_order)}
    transformed_faces = np.array([[mapping[idx] for idx in face] for face in faces])
    return transformed_faces

if __name__ == '__main__':

    smpl = SMPL(path='superSMPL', device='cuda')
    pose = torch.zeros((1,72), device='cuda')
    # generate random pose
    # pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
    body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False)
    colors = np.zeros((body.shape[1], 3))


    renderer = PoseRenderer()
    body_to_render = body[0]
    # swap the y,z axes of the body_to_render 
    body_to_render = body_to_render[:,[0,2,1]]
    # and rotate it 90 degress along the z axis
    body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device='cuda'), body_to_render.T).T

    # out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body.get_device()))
    # ex.imshow(out[0,:,:,:3].cpu().numpy()).show()

    body = body.cpu().numpy()[0]
    faces = smpl.f_all

    # m = trimesh.Trimesh(vertices=body, faces=faces, process=False, validate=False)
    # # export the mesh to an obj file
    # m.export('custom_data/body.obj')
    # imoprt the mesh from the obj file
    m = trimesh.load('custom_data/body_color_init.ply', process=False, validate=False)
    color_init = np.array(m.visual.vertex_colors)[:,:3]
    colored_vertices = {}
    for idx, color in enumerate(color_init):
        if np.all(color == 255):
            continue
        colored_vertices[idx] = color

    # create a 3d plotly scatter plot fo body vertices with vertex 26157 highlited and limits between -1 and 1
    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,1], z=body[:,2], mode='markers',  marker=dict(color='blue', size=1))])
    # fig.add_trace(go.Scatter3d(x=[body[26157,0]], y=[body[26157,1]], z=[body[26157,2]], mode='markers', marker=dict(color='red', size=5)))
    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
    # fig.show()




    # dist = get_distances(body, get_edges(faces))
    # # # dist = compute_distances(body, faces)
    # np.save('custom_data/dijsktra_int.npy', dist)

    body_ply, colors_ply = get_data()
    # rotate body_ply -90 degress along x axis
    body_ply = body_ply[:,[0,2,1]]
    # rotate body_ply 180 degress along y axis
    body_ply[:,2] = -body_ply[:,2]

    # colored_vertices.pop(13833)
    # colored_vertices[1221] = np.array([0, 128, 0]).astype(np.uint8)
    c = find_corresponding_vertices(body_ply, body, list(colored_vertices.keys()))
    # c = np.delete(c, np.where(c == 180))

    dist = np.load('custom_data/dijsktra2.npy')
    ordered_indices = order_mesh_vertices(body, dist, get_edges(faces))
    # reverse ordered_indices
    # ordered_indices = ordered_indices[::-1]

    c = transform_selected_indices(c, ordered_indices)
    body = body[ordered_indices]
    dist = dist[ordered_indices]
    dist = dist[:,ordered_indices]
    faces = transform_faces(faces, ordered_indices)
    colors = interpolate_colors(dist, list(c), list(colored_vertices.values()))
    np.save('custom_data/body_colors_interpolated_unweighted.npy', colors)
 
    #c = [14110, 3192, 25113, 10822, 27513, 25872]
    # generate a gradient grayscale colors for the vertice based ont heir indices
    # colors = np.zeros((body.shape[0], 3))
    # colors[:,0] = np.arange(body.shape[0]) / body.shape[0]
    # colors[:,1] = np.arange(body.shape[0]) / body.shape[0]
    # colors[:,2] = np.arange(body.shape[0]) / body.shape[0]
    # colors = colors * 255
    # colors = colors.astype(np.uint8)


    # weights = dist[0,:]
    # # define a gradient rgb colors for body based on weights that is not just gray
    # colors = np.zeros((body.shape[0], 3))
    # colors[:,0] = weights
    # colors[:,1] = weights
    # colors[:,2] = weights
    # colors = colors * 255
    # colors = colors.astype(np.uint8)

    # # define a gradient rgb colors for body based on their distance from vertex 1221
    # colors = np.zeros((body.shape[0], 3))
    # colors[:,0] = dist[1221,:]
    # colors[:,1] = dist[1221,:]
    # colors[:,2] = dist[1221,:]
    # norm = colors / np.max(colors)
    # colors = colors * 255

    # get the distances between all vertices using dijsktra algorithm without get_dijsktra
    # from scipy.sparse import csr_matrix
    # from scipy.sparse.csgraph import dijkstra



    # define a gradinet color cube for the body that goes from green to red based on the coordinates of the vertices and not their indices
    # colors = np.zeros((body.shape[0], 3))
    # colors[:,0] = (body[:,0] - np.min(body[:,0])) / (np.max(body[:,0]) - np.min(body[:,0]))
    # colors[:,1] = (body[:,1] - np.min(body[:,1])) / (np.max(body[:,1]) - np.min(body[:,1]))
    # colors[:,2] = (body[:,2] - np.min(body[:,2])) / (np.max(body[:,2]) - np.min(body[:,2]))
    fig = go.Figure(data=[go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=colors)])
    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
    fig.show()  
    # np.save('custom_data/pose_colors.npy', colors.astype(np.float32))

    # same as a scatter plot
    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,2], z=body[:,1], mode='markers', marker=dict(size=2, color=colors, opacity=0.8))])
    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
    # fig.show()

    # save body with colors into a ply file
    m = trimesh.Trimesh(vertices=body, faces=faces, vertex_colors=colors)
    m.export('custom_data/body_colors.ply')


    
    pass