import numpy as np
import os
import sys
import pickle
import plotly.express as ex
import plotly.graph_objects as go
import torch
import open3d as o3d
import trimesh
# from vedo import trimesh as trimesh_plot
from scipy.spatial import KDTree
from scipy.sparse import lil_matrix
from scipy.optimize import minimize

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

from smpl.smpl_torch import SMPL

# import mask faces
from custom_utils.custom_utils import _maskFaces

smpl = SMPL(path='smpl')

def energy_function(params, source_points, target_points, indices):
    num_points = source_points.shape[0]
    num_dims = source_points.shape[1]
    
    # Reshape the parameters
    weights = params[:num_points * num_dims].reshape((num_points, num_dims))
    translation = params[num_points * num_dims:num_points * num_dims + num_dims]
    control_points = params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
    
    # Compute the transformed source points
    transformed_points = source_points + np.dot(weights, control_points.T)
    
    # Compute the TPS energy
    tps_energy = 0.5 * np.sum(np.square(target_points[indices] - transformed_points))
    
    return tps_energy

def non_rigid(source_points, target_points):
    tree = KDTree(target_points)
    distances, indices = tree.query(source_points)

    # Initialize the optimization parameters
    num_points = source_points.shape[0]
    num_dims = source_points.shape[1]
    initial_params = np.zeros(num_points * num_dims + num_dims + num_points)

    # Define the optimization bounds
    bounds = [(None, None)] * (num_points * num_dims)
    bounds += [(None, None)] * num_dims
    bounds += [(None, None)] * (num_points * num_dims)

    # Perform the optimization
    result = minimize(energy_function, initial_params, method='L-BFGS-B',
                    args=(source_points, target_points, indices), bounds=bounds)

    # Get the optimized parameters
    optimized_params = result.x

    # Reshape the parameters
    optimized_weights = optimized_params[:num_points * num_dims].reshape((num_points, num_dims))
    optimized_translation = optimized_params[num_points * num_dims:num_points * num_dims + num_dims]
    optimized_control_points = optimized_params[num_points * num_dims + num_dims:].reshape((-1, num_dims))

    transformed_points = source_points + np.dot(optimized_weights, optimized_control_points.T)
    return transformed_points


def register_smpl():
    # generate body with zero pose and shape
    pose = torch.zeros((1, 72))
    shape = torch.zeros((1, 11))
    smpl_out = smpl(pose, shape)[0][0]

    with open('/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl', "rb") as f:
        data = pickle.load(f)
        garment = data['vertices'][100]
        faces = data['faces']

    # rotate garment x degrees along y axis
    # convert degrees to radians
    degrees = 107
    theta = degrees * np.pi / 180
    rot = np.array([[np.cos(theta), 0, np.sin(theta)],
                    [0, 1, 0],
                    [-np.sin(theta), 0, np.cos(theta)]])
    garment = np.dot(rot, garment.T).T

    # rotate garment x degrees along x axis
    # convert degrees to radians
    degrees = 5
    theta = degrees * np.pi / 180
    rot = np.array([[1, 0, 0],
                    [0, np.cos(theta), -np.sin(theta)],
                    [0, np.sin(theta), np.cos(theta)]])
    garment = np.dot(rot, garment.T).T

    # rotate garment x degrees along z axis
    # convert degrees to radians
    degrees = -5
    theta = degrees * np.pi / 180
    rot = np.array([[np.cos(theta), -np.sin(theta), 0],
                    [np.sin(theta), np.cos(theta), 0],
                    [0, 0, 1]])
    garment = np.dot(rot, garment.T).T


    garment[:,0] -= 0.06
    garment[:,1] -= 0.03
    garment[:,2] -= 0.015
    
    # visualize the body and garment
    # fig = go.Figure(data=[go.Mesh3d(x=smpl_out[:,0], y=smpl_out[:,1], z=smpl_out[:,2], i=smpl.f[:,0], j=smpl.f[:,1], k=smpl.f[:,2])])
    # fig.add_trace(go.Mesh3d(x=garment[:,0], y=garment[:,1], z=garment[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2]))
    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
    # fig.show()

    # non_rigid(garment, smpl_out.numpy())

    # Load the source and target meshes
    source_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(garment), o3d.utility.Vector3iVector(faces))
    target_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(smpl_out), o3d.utility.Vector3iVector(smpl.f))
    # flip source and target mesh
    source_mesh, target_mesh = target_mesh, source_mesh

    source_mesh.compute_vertex_normals()
    source_mesh = source_mesh.filter_smooth_laplacian(number_of_iterations=10)
    # plot the smoothed mesh with plotly
    # fig = go.Figure(data=[go.Mesh3d(x=np.asarray(source_mesh.vertices)[:,0], y=np.asarray(source_mesh.vertices)[:,1], z=np.asarray(source_mesh.vertices)[:,2], i=np.asarray(source_mesh.triangles)[:,0], j=np.asarray(source_mesh.triangles)[:,1], k=np.asarray(source_mesh.triangles)[:,2])])
    # fig.show()

    # Perform initial alignment (translation and/or scaling)
    # You can modify this step based on your specific needs
    # For example, you can manually set an initial transformation
    init_transformation = [[1.0, 0.0, 0.0, 0.0],
                        [0.0, 1.0, 0.0, 0.0],
                        [0.0, 0.0, 1.0, 0.0],
                        [0.0, 0.0, 0.0, 1.0]]

    # Create the ICP object and set the parameters
    icp = o3d.pipelines.registration.TransformationEstimationPointToPoint()
    criteria = o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000)
        
    # Perform ICP registration
    reg_result = o3d.pipelines.registration.registration_icp(o3d.geometry.PointCloud(source_mesh.vertices),
                                                             o3d.geometry.PointCloud(target_mesh.vertices),
                                                             0.075,
                                                             init_transformation, icp, criteria)

    # Get the aligned source mesh after registration
    aligned_mesh = source_mesh.transform(reg_result.transformation)

    # Visualize the aligned mesh and target mesh with plotly
    fig = go.Figure(data=[go.Mesh3d(x=np.asarray(aligned_mesh.vertices)[:,0], y=np.asarray(aligned_mesh.vertices)[:,1], z=np.asarray(aligned_mesh.vertices)[:,2], i=np.asarray(aligned_mesh.triangles)[:,0], j=np.asarray(aligned_mesh.triangles)[:,1], k=np.asarray(aligned_mesh.triangles)[:,2])])
    fig.add_trace(go.Mesh3d(x=np.asarray(target_mesh.vertices)[:,0], y=np.asarray(target_mesh.vertices)[:,1], z=np.asarray(target_mesh.vertices)[:,2], i=np.asarray(target_mesh.triangles)[:,0], j=np.asarray(target_mesh.triangles)[:,1], k=np.asarray(target_mesh.triangles)[:,2]))
    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
    fig.show()

    # Get the vertices of the aligned source mesh and target mesh
    aligned_source_vertices = np.asarray(aligned_mesh.vertices)
    target_vertices = np.asarray(target_mesh.vertices)

    # Perform nearest neighbor search to find corresponding vertex indices
    kd_tree = o3d.geometry.KDTreeFlann(target_mesh)
    correspondences = []
    matched = set()
    for i in range(len(aligned_source_vertices)):
        k = 1
        _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
        # while idx[-1] in matched:
        #     k += 1
        #     _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
        # matched.add(idx[-1])
        correspondences.append((i, idx[-1]))

    # correspondences is a list of tuples where each tuple contains the corresponding vertex indices
    # For example, correspondences = [(0, 10), (1, 5), (2, 15), ...]

    # Visualize the correspondences with plotly
    # fig = go.Figure(data=[go.Scatter3d(x=aligned_source_vertices[:,0], y=aligned_source_vertices[:,1], z=aligned_source_vertices[:,2], mode='markers', marker=dict(color='red', size=1))])
    # fig.add_trace(go.Scatter3d(x=target_vertices[:,0], y=target_vertices[:,1], z=target_vertices[:,2], mode='markers', marker=dict(color='blue', size=1)))
    # for corr in correspondences:
    #     fig.add_trace(go.Scatter3d(x=[aligned_source_vertices[corr[0],0], target_vertices[corr[1],0]], y=[aligned_source_vertices[corr[0],1], target_vertices[corr[1],1]], z=[aligned_source_vertices[corr[0],2], target_vertices[corr[1],2]], mode='lines', line=dict(color='green', width=1)))
    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
    # fig.show()

    registered = np.zeros((14475, 3))
    mask = np.zeros((14475), dtype=bool)
    for i in range(len(correspondences)):
        registered[correspondences[i][1]] = aligned_source_vertices[correspondences[i][0]]
        mask[correspondences[i][1]] = True

    # F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)

    # registered_masked, masked_faces = _maskFaces(registered, np.asarray(target_mesh.triangles), mask)
    registered_masked, masked_faces = registered, np.asarray(target_mesh.triangles)

    # plot aligned source mesh with the target topology
    fig = go.Figure(data=[go.Mesh3d(x=registered_masked[:,0], y=registered_masked[:,1], z=registered_masked[:,2], i=masked_faces[:,0], j=masked_faces[:,1], k=masked_faces[:,2])])
    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
    fig.show()

    return smpl


def convert_vto(folder, out_folder):
    os.makedirs(out_folder, exist_ok=True)
    os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
    os.makedirs(os.path.join(out_folder, 'mesh_posed_body_no_orient'), exist_ok=True)
    os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
    os.makedirs(os.path.join(out_folder, 'pose_shape'), exist_ok=True)
    F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)

    files = os.listdir(folder)
    # files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl']
    for file in files:
        # read pickle files and get vetices and faces from them
        with open(os.path.join(folder, file), "rb") as f:
            data = pickle.load(f)
            vertices = data['vertices']
            faces = data['faces']
            pose = data['pose']
            shape = data['shape']
            trans = data['translation']
            id = file.split('.')[0]
            # pose[:, :3] = 0
            id = id.replace('_', '-')
            subject = id.split('-')[2]
            sequence = '_'.join(id.split('-')[2:])
            # cmu_data = np.load(f'/data-net/code/vto-garment-collisions/assets/CMU/{subject}/{sequence}_poses.npz')
            # shape = np.tile(cmu_data['betas'], (shape.shape[0], 1))
            # pose = cmu_data['poses']

            body = smpl(torch.tensor(pose), torch.tensor(np.concatenate((shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)))[0].numpy()
            pose_shape = np.concatenate((pose, shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)
            
            # for i in range(vertices.shape[0]):
            #     np.savez_compressed(os.path.join(out_folder, 'mesh_posed_body_no_orient', f'{id}_{i}.npz'), mesh_posed_body_no_orient=body[i])
            #     np.save(os.path.join(out_folder, 'pose_shape', f'{id}_{i}.npy'), pose_shape[i])


            # plot mesh with plotly
            frame = 0
            fig = go.Figure(data=[go.Mesh3d(x=vertices[frame,:,0], y=vertices[frame,:,1], z=vertices[frame,:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])])
            fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,1], z=body[frame,:,2], i=F[:,0], j=F[:,1], k=F[:,2]))
            fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
            fig.show()

            
            pass
    pass


if __name__ == '__main__':
    # register_smpl()
    convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted')
    pass