import os
import sys
import shutil
import tqdm
import random
import math
import collections
import tps
import numpy as np
import scipy as sp
import trimesh as tm
import cmapy as cm
import scipy.io as sio
import matplotlib.pyplot as plt
import plotly.express as ex



def m2uv_array(m2uv, is_smpl_orig=False):
    result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
    for k, v in m2uv.items():
        if len(v) == 2:
            result[k] = list(v)
        elif len(v) == 1:
            x = v.pop()
            result[k] = [x, x]
        else:
            assert False

    return result


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 plot_3d(x, y=None):
    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlim3d(-1, 1)
    ax.set_ylim3d(-1, 1)
    ax.set_zlim3d(-1, 1)
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
    plt.show()

def plot_uv(x, y=None):
    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(111)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.scatter(x[:, 0], x[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    plt.show()

""" 
Mesh to UV map
Computes correspondences between 3D mesh and UV map
NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
"""
def mesh2UV(F, Ft):
	m2uv = {v: set() for f in F for v in f}
	for f, ft in zip(F, Ft):
		for v, vt in zip(f, ft):
			m2uv[v].add(vt)
	# m2uv = {k:list(v) for k,v in m2uv.items()}
	return m2uv


def filter_m2u(m2u, to_keep):
    for key in m2u:
        if len(m2u[key]) == 2:
            v = list(m2u[key])
            if v[0] in to_keep:
                m2u[key].remove(v[1])
            else:
                m2u[key].remove(v[0])
    return m2u


def convert_to_numbers(m2u):
    for key in m2u:
        m2u[key] = m2u[key].pop()
    return m2u


def get_masked_index(mask, index):
    m = np.zeros((mask.shape[0]), dtype=bool)
    m[index] = True
    m = m[mask]
    new_index = np.argmax(m)
    return new_index


def get_closest_point(verts_1, verts_2):
    dist = sp.spatial.distance.cdist(verts_1, verts_2)
    closest = np.argmin(dist, axis=-1)
    return closest

def get_normal_edge(p1, p2):
    d = p2 - p1
    return np.array([-d[1], d[0]])

def get_normal_vertex(p1, p2, p3):
    n1 = get_normal_edge(p1, p2)
    n2 = get_normal_edge(p2, p3)
    n = (n1 + n2)
    n = n / np.linalg.norm(n) / 20
    return n


def get_distance(p1, p2, points):
    d = np.array([np.linalg.norm(np.cross(p2 - p1, point - p1)) / np.linalg.norm(p2 - p1) for point in points])
    return d


def get_normals(verts):
    v = verts.copy()
    v = np.concatenate([v[-1][np.newaxis, :], v, v[0][np.newaxis, :]])
    normals = np.zeros_like(verts)
    for i in range(1, v.shape[0] - 1):
        normals[i - 1] = get_normal_vertex(v[i - 1], v[i], v[i + 1])
    return normals


def get_closest_points_normals(verts1, verts2):
    dist_2 = sp.spatial.distance.cdist(verts1, verts2)
    normals = get_normals(verts1)

    # fig = plt.figure(figsize=(15, 15))
    # ax = fig.add_subplot(111)
    # ax.set_xlim(0, 1)
    # ax.set_ylim(0, 1)
    # ax.set_aspect('equal')
    # ax.scatter(verts1[:, 0], verts1[:, 1])
    # for p1, p2 in zip(verts1, verts1 + normals):
    #     ax.plot([p1[0], p2[0]], [p1[1], p2[1]])
    # plt.show()

    closest = []
    i = 0
    for p1, p2 in zip(verts1, verts1 + normals):
        dist = get_distance(p1, p2, verts2)
        closest.append(np.argmin(dist + dist_2[i, :]))
        i += 1

    return np.array(closest)


def back():
    # load data
    data_body = np.load("Blender/uv_data_1.npz", allow_pickle=True)
    data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
    
    # translate back uv to match body
    uv_back = data_back.f.Vt
    # uv_back[:, 0] = uv_back[:, 0] + 0.1
    # uv_back[:, 1] = uv_back[:, 1] + 0.42

    # find front and back indices in uv
    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)

    # get mesh to uv mapping for the back vertices
    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
    m2u_body = filter_m2u(m2u_body, set(back))
    m2u_body = convert_to_numbers(m2u_body)

    m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
    m2u_back = convert_to_numbers(m2u_back)

    # create mask for filtering vertices belonging to the back
    body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
    body_back_mask[back] = True

    # define which points should match after the transform on the mesh
    matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
    
    # map matching point to uv coordinates
    for i in range(matching_points.shape[0]):
        matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]

    x = data_body.f.Vt[body_back_mask, :]
    y = uv_back

    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(121)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(x[:, 0], x[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    for pair in matching_points:
        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)

    # get transformation from body to back contour
    source =  x[matching_points[:, 0], :]
    destination = y[matching_points[:, 1], :]
    trans = tps.TPS(source, destination)
    transformed = trans(data_body.f.Vt[back, :])

    ax = fig.add_subplot(122)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(transformed[:, 0], transformed[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    for pair in matching_points:
        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)

    plt.show()

    data_back_high = np.load("Blender/uv_data_plane_1_high_res.npz", allow_pickle=True)
    y_high = data_back_high.f.Vt
    # y_high[:, 0] = y_high[:, 0] + 0.1
    # y_high[:, 1] = y_high[:, 1] + 0.42
    
    closest_points = get_closest_points_normals(y_high, transformed)
    # closest_points = get_closest_point(y_high, transformed)
    unique_mask = np.zeros_like(closest_points, dtype=bool)
    unique_mask[np.unique(closest_points, return_index=True)[1]] = True
    matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
    # matching_points_2 = []
    # for i, p in enumerate(matching_points):
    #     if i % 10 == 0:
    #         matching_points_2.append(p)
    # matching_points = np.array(matching_points_2)

    source_2 = transformed[matching_points[:, 0], :]
    destination_2 = y_high[matching_points[:, 1], :]
    trans_2 = tps.TPS(source_2, destination_2)
    transformed_2 = trans_2(transformed)



    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(121)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(source_2[:, 0], source_2[:, 1])
    ax.scatter(destination_2[:, 0], destination_2[:, 1])
    for i, pair in enumerate(zip(source_2, destination_2)):
        if i % 5 == 0:
            s, d = pair
            ax.scatter([s[0], d[0]], [s[1], d[1]])

    ax = fig.add_subplot(122)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
    ax.scatter(y_high[:, 0], y_high[:, 1])
    plt.show()

    return transformed_2


def match_back_to_front(front_uv):
    # load data
    data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
    # data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
    
    # translate back uv to match body
    # uv_back = data_back.f.Vt
    # uv_back[:, 0] = uv_back[:, 0] + 0.1
    # uv_back[:, 1] = uv_back[:, 1] + 0.42

    # find front and back indices in uv
    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)

    # get mesh to uv mapping for the back vertices
    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
    m2u_body = filter_m2u(m2u_body, set(back))
    m2u_body = convert_to_numbers(m2u_body)

    # m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
    # m2u_back = convert_to_numbers(m2u_back)

    # create mask for filtering vertices belonging to the back
    body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
    body_back_mask[back] = True
    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
    body_front_mask[front] = True

    # define which points should match after the transform on the mesh
    # matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
    
    # # map matching point to uv coordinates
    # for i in range(matching_points.shape[0]):
    #     matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]

    m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
    m2u_body_all = m2uv_array(m2u_body_all)
    border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
    border_pairs = m2u_body_all[border_pairs, :]
    matching_points = border_pairs
    matching_points_sorted = []
    for p in matching_points:
        if p[0] in front:
            matching_points_sorted.append([p[0], p[1]])
        elif p[0] in back:
            matching_points_sorted.append([p[1], p[0]])
        else:
            assert False
    # add neck V shape control point
    matching_points_sorted.append([m2u_body_all[10228][0], m2u_body_all[33][0]])
    matching_points = np.array(matching_points_sorted)
    mm = matching_points.copy()

    for i in range(matching_points.shape[0]):
        matching_points[i, :] = [get_masked_index(body_back_mask, matching_points[i,1]), get_masked_index(body_front_mask, matching_points[i,0])]

    temp = np.zeros_like(data_body.f.Vt)
    temp[front, :] = front_uv

    x = data_body.f.Vt[body_back_mask, :]
    y = temp[body_front_mask, :] # front_uv # 

    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(121)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(x[:, 0], x[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    for pair in matching_points:
        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)

    # get transformation from body to back contour
    source =  x[matching_points[:, 0], :]
    destination = y[matching_points[:, 1], :]
    trans = tps.TPS(source, destination)
    transformed = trans(data_body.f.Vt[back, :])

    ax = fig.add_subplot(122)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(transformed[:, 0], transformed[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    for pair in matching_points:
        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)

    plt.show()


    return transformed


def generate_mask():
    import cv2
    res = 1024
    front = np.load("Blender/uv_data_plane_0_high_res_2.npz", allow_pickle=True)
    back = np.load("Blender/uv_data_plane_1_high_res_2.npz", allow_pickle=True)
    points = np.round(np.concatenate((front.f.Vt, back.f.Vt)) * res).astype(np.int32)
    img = np.zeros((res, res, 3))
    img[points[:, 0], points[:, 1], :] = 255
    cv2.imwrite('MeshUtils/mask_1024.png', img)

    # do paint now

    img = cv2.imread('MeshUtils/mask_1024_2.png')
    img = img[:, :, 0] == 255
    np.save('MeshUtils/mask_1024.npy', img)

def generate_mask_2():
    import cv2
    res = 256
    front = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
    back = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
    points = np.floor(np.concatenate((front.f.Vt, back.f.Vt + [1, 0])) * res).astype(np.int32)
    img = np.zeros((res * 2, res, 3))
    img[points[:, 0], points[:, 1], :] = 255
    cv2.imwrite('MeshUtils/mask_256x512.png', img)

    # do paint now

    img = cv2.imread('MeshUtils/mask_256x512.png')
    img = img[:, :, 0] > 240
    np.save('MeshUtils/mask_256x512.npy', img)

    pass


def generate_uv_mask_border():
    import cv2
    mask = np.load('MeshUtils/mask_256x512.npy')
    uv_res = 256
    mask_border = (cv2.resize(mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
    mask_border = cv2.erode(mask_border.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)

    mask[mask_border == True] = 0
    # plt.imshow(mask)
    idx_front = np.transpose(np.vstack(np.nonzero(mask[:256, :])), [1,0])

    test = np.zeros_like(mask)
    test[idx_front[:,0], idx_front[:,1]] = 1
    idx_back = idx_front.copy()
    idx_back[:, 0] += 256
    test[idx_back[:,0], idx_back[:,1]] = 1


    pass



def front():
    # load data
    data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
    data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
    
    # translate back uv to match body
    uv_front = data_front.f.Vt
    # uv_back[:, 0] = uv_back[:, 0] + 0.1
    # uv_back[:, 1] = uv_back[:, 1] + 0.42

    # find front and back indices in uv
    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)

    # get mesh to uv mapping for the back vertices
    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
    m2u_body = filter_m2u(m2u_body, set(front))
    m2u_body = convert_to_numbers(m2u_body)

    m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
    m2u_front = convert_to_numbers(m2u_front)

    # create mask for filtering vertices belonging to the back
    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
    body_front_mask[front] = True

    # define which points should match after the transform on the mesh
    matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
    # matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4793, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8450, 14], [13566, 18], [11767, 19], [8668, 20], [5012, 21], [12716, 22], [10917, 23], [8280, 24], [4617, 25], [2760, 26], [911, 27] ])

    # map matching point to uv coordinates
    for i in range(matching_points.shape[0]):
        matching_points[i, :] = [get_masked_index(body_front_mask, m2u_body[matching_points[i,0]]), m2u_front[matching_points[i, 1]]]

    x = data_body.f.Vt[body_front_mask, :]
    y = uv_front

    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(121)
    # ax.set_xlim(0, 1)
    # ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(x[:, 0], x[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    for pair in matching_points:
        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)

    # get transformation from body to front contour
    source = x[matching_points[:, 0], :]
    destination = y[matching_points[:, 1], :]
    trans = tps.TPS(source, destination)
    transformed = trans(data_body.f.Vt[front, :])

    # transformed = (transformed - transformed.min()) / (transformed.max() - transformed.min())

    ax = fig.add_subplot(122)
    # ax.set_xlim(0, 1)
    # ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(transformed[:, 0], transformed[:, 1])
    if y is not None:
        ax.scatter(y[:, 0], y[:, 1])
    for pair in matching_points:
        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)

    plt.show()

    data_front_high = np.load("custom_data/uv_map/uv_data_plane_stretched_front_high.npz", allow_pickle=True)
    y_high = data_front_high.f.Vt
    # y_high[:, 0] = y_high[:, 0] + 0.1
    # y_high[:, 1] = y_high[:, 1] + 0.42
    
    closest_points = get_closest_points_normals(y_high, transformed)
    # closest_points = get_closest_point(y_high, transformed)
    unique_mask = np.zeros_like(closest_points, dtype=bool)
    unique_mask[np.unique(closest_points, return_index=True)[1]] = True
    matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
    # matching_points_2 = []
    # for i, p in enumerate(matching_points):
    #     if i % 10 == 0:
    #         matching_points_2.append(p)
    # matching_points = np.array(matching_points_2)

    source_2 = transformed[matching_points[:, 0], :]
    destination_2 = y_high[matching_points[:, 1], :]
    trans_2 = tps.TPS(source_2, destination_2)
    transformed_2 = trans_2(transformed)



    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(121)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(source_2[:, 0], source_2[:, 1])
    ax.scatter(destination_2[:, 0], destination_2[:, 1])
    for i, pair in enumerate(zip(source_2, destination_2)):
        if i % 5 == 0:
            s, d = pair
            ax.scatter([s[0], d[0]], [s[1], d[1]])

    ax = fig.add_subplot(122)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_aspect('equal')
    ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
    ax.scatter(y_high[:, 0], y_high[:, 1])
    plt.show()

    return transformed_2, mesh2UV(data_body.f.F, data_body.f.Ft), front, back


def smoothing():
    data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
    uv_map = np.load('MeshUtils\stretched_uv_map_3.npy')
    verts = np.concatenate((uv_map, np.ones((15068, 1))), axis=-1)
    m = tm.Trimesh(verts, data_body.f.Ft, process=False, validate=False)


    m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
    m2u_body_all = m2uv_array(m2u_body_all)
    border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
    border_pairs = m2u_body_all[border_pairs, :]
    pinned_vertices = border_pairs.reshape(-1)

    lap = tm.smoothing.laplacian_calculation(m, equal_weight=False, pinned_vertices=pinned_vertices)
    # tm.smoothing.filter_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
    # tm.smoothing.filter_mut_dif_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
    # tm.smoothing.filter_humphrey(m, iterations=100, laplacian_operator=lap)
    tm.smoothing.filter_taubin(m, iterations=100, laplacian_operator=lap)


    x = uv_map
    y = m.vertices[:, :2]

    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(211)
    ax.set_xlim(0, 2)
    ax.set_ylim(0, 1)
    # ax.set_aspect('equal')
    ax.scatter(x[:, 0], x[:, 1])

    ax = fig.add_subplot(212)
    ax.set_xlim(0, 2)
    ax.set_ylim(0, 1)
    # ax.set_aspect('equal')
    ax.scatter(y[:, 0], y[:, 1])
    plt.show()


    pass



if __name__ == '__main__':
    # generate_mask_2()
    # smoothing()
    # generate_uv_mask_border()

    front_uv, m2uv, front_mask, back_mask = front()
    back_uv = match_back_to_front(front_uv)
    # back_uv = back()

    final = np.zeros((15068, 2))
    final[front_mask] = front_uv
    final[back_mask] = back_uv + [1, 0]
    np.save('MeshUtils\stretched_uv_map_3.npy', final)

    pass