import os
import sys
import random
import numpy as np
import plotly.graph_objects as go
import torch
from torch.nn.functional import mse_loss
from pytorch3d.structures import Meshes
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
from pymeshfix import MeshFix
import scipy.spatial as spatial

os.environ["CUDA_LAUNCH_BLOCKING"] = '1'

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import (
    read_garment_mask,
    read_garment_template,
    read_posed_garment_frame,
    loadInfo,
    read_posed_body_frame,
    read_posed_body,
    read_posed_garment,
    read_curvatures,
    read_curvature_frame,
    writeOBJ,
    readOBJ,
    quads2tris
)
from smpl.smpl_torch import SMPL
# from custom_utils.custom_utils import _maskFaces

device = 'cuda'
smpl = SMPL(device, 'smpl')


import torch.nn.functional as F
from pytorch3d.structures import Meshes

def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
    """
    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
    find the closest vertices from mesh2 to these vertices.

    Args:
        mesh1: A PyTorch3D Meshes object representing the first mesh.
        mesh2: A PyTorch3D Meshes object representing the second mesh.
        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.

    Returns:
        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
    """
    # Extract the vertices of mesh1 and mesh2.
    vertices1 = mesh1.verts_packed()
    vertices2 = mesh2.verts_packed()

    # Extract the vertices of mesh1 that correspond to the given indices.
    vertices1_selected = vertices1[vertex_indices]

    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)

    # Find the indices of the closest vertices in mesh2.
    closest_indices = torch.argmin(distances, dim=1)

    return closest_indices


def _maskFaces(V, F, mask, return_unfixed=False):
    V_mask = mask
    F_mask = np.take(V_mask, F.reshape(-1), axis=0)
    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3

    V_ind = np.zeros(V_mask.shape)
    V_ind[V_mask] = np.arange(V_mask.sum())
    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
    V_new = V[V_mask, :]

    if return_unfixed:
        return V_new, F_new, F[F_mask, :]
    else:
        return V_new, F_new

def get_frames(folder, frames_file):
    frames = []
    with open(os.path.join(folder, frames_file), 'r') as f:
        for line in f:
            frames.append(line.strip())
    return frames


def plot(mesh, faces):
    fig = go.Figure()

    fig.add_trace(go.Mesh3d(
                x=mesh[:,0],
                y=mesh[:,1],
                z=mesh[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),)
    
    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
                    )
    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
                yaxis = dict(nticks=4, range=[-1,1],),
                zaxis = dict(nticks=4, range=[-1,1],),)
    fig.update_layout(scene = helper,
                        )

    fig.show()


def build_lookup_table(repeated):
    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element

    lookup = {}
    for i in range(len(repeated)):
        for j in range(len(repeated[i])):
            if len(repeated[i]) == 1:
                continue
            lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
    return lookup


from data.cloth3d_io import readPC2Frame
import scipy.io as sio
import scipy
import scipy.spatial as spatial

def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
    pc16_path = os.path.join(folder, id, garment + '.pc16')
    verts = readPC2Frame(pc16_path, frame, True)

    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))

    if deorient:
        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
        verts = R.apply(verts)

    return verts, map

def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
    verts = readPC2Frame(pc16_path, frame, True)

    if deorient:
        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
        verts = R.apply(verts)

    return verts

def edge_loss_custom(x, y, edges):
    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
    x_e = torch.sum(x_e ** 2, -1)
    x_e = torch.sqrt(x_e + 1e-9)
    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
    y_e = torch.sum(y_e ** 2, -1)
    y_e = torch.sqrt(y_e + 1e-9)
    d_e = y_e - x_e
    loss = torch.mean(d_e ** 2, dim=-1)[0]
    return loss


def crotch(mesh, body, crotch_idx, nr=75):
    crotch_middle_on_body = crotch_idx
    # create kdtree for body and mesh
    body_tree = spatial.cKDTree(body)
    mesh_tree = spatial.cKDTree(mesh)
    # find closest point on mesh from crotch middle on body along the z axis

    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)

    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]

    body_verts = body[crotch_body]
    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
    mesh_verts = spatial.cKDTree(mesh_verts)
    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
    crotch_mesh = crotch_mesh[closest_mesh_verts]

    return crotch_body, crotch_mesh


def find_closest_point_along_z(point, verts, eps=0.01):
    min_dist = np.inf
    closest_point = None
    for i, v in enumerate(verts):
        dist = np.linalg.norm(point - v)
        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
            min_dist = dist
            closest_point = i
    return closest_point


def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
    """
    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.

    Args:
        mesh: A PyTorch3D Meshes object representing the mesh.

    Returns:
        A tensor of shape (E,) representing the lengths of all edges in the mesh,
        where E is the total number of edges in the mesh.
    """
    # Extract the vertices and faces of the mesh.
    vertices = mesh.verts_packed()
    faces = mesh.faces_packed()

    # Compute the pairwise distances between all pairs of vertices that form an edge.
    edge_indices = faces.view(-1, 3)
    edge_vertices = vertices[edge_indices]
    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)

    return edge_lengths

# try to optimize the already registered garment ??????????

def register(folder, frame):
    id, frame_nr, garment = frame.split()
    frame_nr = int(frame_nr)
    info = loadInfo(os.path.join(folder, id, 'info.mat'))
    raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
    # convert it to indices
    mask = raw_mask > 0

    _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
    faces = quads2tris(faces)
    template = np.load('temp/mapping_cloth.npy')
    mapping_mask = np.load('temp/mapping_mask.npy')
    mapping_faces = np.load('temp/mapping_faces.npy')
    mapping_weights = np.load('temp/mapping_weights.npy')

    mapping_faces, mapping_weights = mapping_faces[mapping_mask], mapping_weights[mapping_mask]

    is_skirt = (garment == 'Dress' or garment == 'Skirt')
    skirt_mod = '_skirt' if is_skirt else ''
    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
    _, registered_faces = _maskFaces(np.zeros((14475, 3)), F, mapping_mask)


    template = torch.tensor(template).to(device).to(torch.float32)
    template_orig = template.clone()
    faces = torch.tensor(faces).to(device)
    mapping_faces = torch.tensor(mapping_faces).to(device)
    mapping_weights = torch.tensor(mapping_weights).to(device).to(torch.float32)
    F = torch.tensor(F).to(device)
    registered_faces = torch.tensor(registered_faces).to(device)

    # verts = template.cpu().numpy()
    # star_body = star_body.cpu().numpy()
    # faces = faces.detach().cpu().numpy()
    # F = F.detach().cpu().numpy()
    # fig = go.Figure(data=[
    #     go.Mesh3d(
    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
    #         color='lightpink', opacity=0.50),
    #     go.Mesh3d(
    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
    #         i=F[:,0], j=F[:,1], k=F[:,2],
    #         color='lightblue', opacity=0.50),
    #     go.Scatter3d(
    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
    #         mode='markers', marker=dict(size=2, color='red')),
    #     go.Scatter3d(
    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
    #         mode='markers', marker=dict(size=2, color='blue')),
    # ])
    # # make it cuve and limit to -1 and -1
    # fig.update_layout(scene = dict(
    #     xaxis = dict(range=[-1,1],),
    #     yaxis = dict(range=[-1,1],),
    #     zaxis = dict(range=[-1,1],),),)
    # fig.show()

    mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
    avg_edge_length = mesh_edge_lengths(mesh2).mean()

    # Define the meshes to optimize
    registered = (template[mapping_faces,:] * mapping_weights[:,:,None]).sum(1)

    optim_mask = torch.ones(registered.shape[0]).to(device).to(torch.bool)
    # registered_tree = spatial.KDTree(template.detach().cpu().numpy())
    # for i in range(registered.shape[0]):
    #     # if distance between registerd and template is too big set it to require grad that weight
    #     res = registered_tree.query(registered[i].detach().cpu().numpy(), k=1)
    #     if res[0] > avg_edge_length / 4:
    #         optim_mask[i] = True

    mapping_weights_masked = mapping_weights[optim_mask]
    mapping_weights_nonmasked = mapping_weights[~optim_mask]
    mapping_weights_masked.requires_grad = True

    # mesh1 = Meshes(verts=registered[None,:,:], faces=registered_faces[None,:,:])


    # Define the optimizer and learning rate
    optimizer = torch.optim.Adam([mapping_weights_masked], lr=0.001)
    # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)

    # Define the loss functions
    def laplacian_smoothing_loss(mesh):
        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot

    def edge_loss(mesh):
        return mesh_edge_loss(mesh, target_length=avg_edge_length)

    def normal_consistency_loss(mesh):
        return mesh_normal_consistency(mesh)

    # Optimization loop
    for i in range(4000):
        optimizer.zero_grad()

        # find boundaries with the find holes function?
        registered_masked = (template[mapping_faces[optim_mask],:] * mapping_weights_masked[:,:,None]).sum(1)
        registered_nonmasked = (template[mapping_faces[~optim_mask],:] * mapping_weights_nonmasked[:,:,None]).sum(1)
        # registered = torch.cat((registered_masked, registered_nonmasked), 0)
        registered = torch.zeros_like(mapping_weights)
        registered[optim_mask] = registered_masked
        registered[~optim_mask] = registered_nonmasked

        # Forward pass
        mesh1_out = Meshes(verts=registered[None,:,:], faces=registered_faces[None,:,:])
        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()),
                                        )
        loss: torch.Tensor = ch_loss * 20 #+ ch_norm_loss * 10
        lap_loss = laplacian_smoothing_loss(mesh1_out)
        e_loss = edge_loss(mesh1_out)
        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
        norm_loss = normal_consistency_loss(mesh1_out)
        loss += lap_loss * 0.05 + e_loss * 20 + norm_loss * 0.1

        # Backward pass
        loss.backward()
        optimizer.step()

        # Print the loss
        if i % 10 == 0:
            print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
    
    # create a plotly plot with the body and the mesh1_out
    registered = mesh1_out.verts_packed().detach().cpu().numpy()
    template = mesh2.verts_packed().detach().cpu().numpy()
    faces = faces.detach().cpu().numpy()
    F = F.detach().cpu().numpy()
    optim_mask = optim_mask.detach().cpu().numpy()
    registered_faces = registered_faces.detach().cpu().numpy()
    # faces, F = F, faces
    # verts, star_body = star_body, verts

    fig = go.Figure(data=[
        go.Mesh3d(
            x=template[:,0], y=template[:,1], z=template[:,2],
            i=faces[:,0], j=faces[:,1], k=faces[:,2],
            color='lightblue', opacity=0.80),
        go.Mesh3d(
            x=registered[:,0], y=registered[:,1], z=registered[:,2],
            i=registered_faces[:,0], j=registered_faces[:,1], k=registered_faces[:,2],
            color='lightpink', opacity=0.50),   
        # go.Scatter3d(
        #     x=template[:,0], y=template[:,1], z=template[:,2],
        #     mode='markers', marker=dict(size=2, color='red')),
        # go.Scatter3d(
        #     x=registered[:,0], y=registered[:,1], z=registered[:,2],
        #     mode='markers', marker=dict(size=2, color='blue')),
    ])
    # make it cuve and limit to -1 and -1
    fig.update_layout(scene = dict(
        xaxis = dict(range=[-1,1],),
        yaxis = dict(range=[-1,1],),
        zaxis = dict(range=[-1,1],),),)
    fig.update_layout(scene_aspectmode='cube')
    fig.show()


    fig = go.Figure(data=[
        go.Scatter3d(
            x=registered[:,0], y=registered[:,1], z=registered[:,2],
            mode='markers', marker=dict(size=4, color='red')),
        go.Scatter3d(
            x=template[:,0], y=template[:,1], z=template[:,2],
            mode='markers', marker=dict(size=2, color='blue')),
    ])
    # make it cuve and limit to -1 and -1
    fig.update_layout(scene = dict(
        xaxis = dict(range=[-1,1],),
        yaxis = dict(range=[-1,1],),
        zaxis = dict(range=[-1,1],),),)
    fig.update_layout(scene_aspectmode='cube')
    fig.show()

    
    # fig = go.Figure(data=[
    #     go.Scatter3d(
    #         x=registered[:optim_mask.sum(),0], y=registered[:optim_mask.sum(),1], z=registered[:optim_mask.sum(),2],
    #         mode='markers', marker=dict(size=4, color='red')),
    #     go.Scatter3d(
    #         x=registered[:,0], y=registered[:,1], z=registered[:,2],
    #         mode='markers', marker=dict(size=2, color='blue')),
    # ])
    # # make it cuve and limit to -1 and -1
    # fig.update_layout(scene = dict(
    #     xaxis = dict(range=[-1,1],),
    #     yaxis = dict(range=[-1,1],),
    #     zaxis = dict(range=[-1,1],),),)
    # fig.update_layout(scene_aspectmode='cube')
    # fig.show()

    # create a plotly wireframe plot of registered
    # fig = go.Figure()
    # for edge in mesh1_out.edges_packed().detach().cpu().numpy():
    #     fig.add_trace(go.Scatter3d(
    #         x=registered[edge,0], y=registered[edge,1], z=registered[edge,2],
    #         mode='lines', line=dict(color='red', width=1)))
    # fig.update_layout(scene = dict(
    #     xaxis = dict(range=[-1,1],),
    #     yaxis = dict(range=[-1,1],),
    #     zaxis = dict(range=[-1,1],),),)
    # fig.update_layout(scene_aspectmode='cube')
    # fig.show()


    pass

if __name__ == "__main__":
    folder = '/data-net/code/cloth3d/train'
    # folder = r'D:\Work\code\cloth3d\train'
    # frames = [f'00761_{i}' for i in range(300)]
    # folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
    frames = get_frames(folder, 'frames_all.txt')
    root = '/'.join(folder.split('/')[:-1])
    type = folder.split('/')[-1]

    use_no_orient = False
    no_shape = False
    is_skirt = False

    skirt_modifier = '_skirt' if is_skirt else ''
    orient_mod = '_no_orient' if use_no_orient else ''
    shape_mod = '_no_shape' if no_shape else ''

    # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
    # faces = np.load('/data-net/dataset/00761/faces/00761.npy')

    for frame in frames:
        frame = random.choice(frames)
        # frame = '00761 0 Tshirt'
        # frame = '01893 256 Jumpsuit'
        frame = '05060 1 Jumpsuit'
        

        register(folder, frame)
        break
        pass