import os
from socket import close
import sys
import tqdm
import math
import time
import random
import shutil
import collections
import argparse
import threading
import numpy as np
import scipy as sp
import trimesh as tm
import cmapy as cm
import cv2
import plotly.express as ex
import plotly.graph_objects as go


def m2uv_array(m2uv):
    result = np.zeros((14475, 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


if __name__ == '__main__':
    F = np.load("custom_utils/smpl_faces.npy")
    uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
    m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
    uv_mask = np.load('custom_utils/mask_256x512.npy')
    
    left = [10602, 14090, 10601, 14299, 10623, 14102, 8739, 13152, 8738, 14292, 10737, 14244, 10738, 14288, 10736, 14243, 8728, 13149, 8730, 13982, 10419, 14241, 8721, 13147, 8723, 14239, 8717, 13145, 8716, 14237, 8711, 13143, 8710, 14235, 10734, 14283, 10729, 14226, 10381, 13948, 10383, 14225, 10725, 14222, 8620, 13107, 8622, 14305, 10761, 14293, 10732, 14233, 10720, 14214, 10721, 14217, 7114]
    right = list(reversed([12418, 7113, 12415, 7112, 12434, 7125, 12494, 7155, 12506, 4964, 11308, 4963, 12423, 7118, 12426, 6752, 12149, 6751, 12427, 7122, 12484, 7127, 12436, 5054, 11344, 5056, 12438, 5060, 11346, 5062, 12440, 5067, 11348, 5066, 12442, 6788, 12183, 5074, 11350, 5072, 12444, 7129, 12489, 7131, 12445, 7130, 12493, 5082, 11353, 5084, 12303, 7009, 12500, 6987, 12291, 6989]))
    corners = np.array([[0,0],[0,1],[2,0],[2,1]])
    
    # left_orig = []
    # for l in left:
    #     left_orig.append(list(m2uv[l])[0])
    #     left_orig.append(list(m2uv[l])[1])
    # ex.scatter(x=np.concatenate((uv_map[left_orig, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[left_orig, 1], corners[:,1]), axis=-1)).show()


    # left_1 = [[0,0],[0,1],[2,0],[2,1]]
    # right_1 = [[0,0],[0,1],[2,0],[2,1]]
    # left_2 = [[0,0],[0,1],[2,0],[2,1]]
    # right_2 = [[0,0],[0,1],[2,0],[2,1]]
    left_1 =  []
    right_1 = []
    left_2 =  []
    right_2 = []
    for l in left:
        print(list(m2uv[l])[0])
        if uv_map[list(m2uv[l])[0], 0] < 1:
            left_1.append(list(m2uv[l])[0])
        else:
            left_2.append(list(m2uv[l])[0])
        if uv_map[list(m2uv[l])[1], 0] < 1:
            left_1.append(list(m2uv[l])[1])
        else:
            left_2.append(list(m2uv[l])[1])
    left_1 = np.array(left_1)[:56]
    left_2 = np.array(left_2)[:56]

    # ex.scatter(x=uv_map[:, 0], y=uv_map[:, 1]).show()

    for r in right:
        print(list(m2uv[r])[0])
        if uv_map[list(m2uv[r])[0], 0] < 1:
            right_1.append(list(m2uv[r])[0])
        else:
            right_2.append(list(m2uv[r])[0])
        if uv_map[list(m2uv[r])[1], 0] < 1:
            right_1.append(list(m2uv[r])[1])
        else:
            right_2.append(list(m2uv[r])[1])
    right_1 = np.array(right_1)
    right_2 = np.array(right_2)

    c = np.ones((60, 3), dtype=np.int32) * 50
    c[5, :] = 200
    # ex.scatter(x=np.concatenate((uv_map[left_1, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[left_1, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
    # ex.scatter(x=np.concatenate((uv_map[right_1, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[right_1, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
    # # ex.scatter(x=uv_map[right_1, 0], y=uv_map[right_1, 1], color=c[:,0]).show()

    # ex.scatter(x=np.concatenate((uv_map[left_2, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[left_2, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
    # ex.scatter(x=np.concatenate((uv_map[right_2, 0], corners[:,0]), axis=-1), y=np.concatenate((uv_map[right_2, 1], corners[:,1]), axis=-1), color=c[:,0]).show()
    # # ex.scatter(x=uv_map[left_2, 0], y=uv_map[left_2, 1]).show()
    # # ex.scatter(x=uv_map[right_2, 0], y=uv_map[right_2, 1]).show()

    # new_pairs_1 = np.concatenate((left_1, right_1), axis=-1)
    # new_pairs_2 = np.concatenate((left_2, right_2), axis=-1)

    # ex.scatter(x=uv_map[new_pairs_1, 0], y=uv_map[new_pairs_1, 1]).show()
    # ex.scatter(x=uv_map[new_pairs_2, 0], y=uv_map[new_pairs_2, 1]).show()


    for i, l in enumerate(left[:56]):
        m2uv[l] = {left_2[i], right_2[i]}
    for i, r in enumerate(right):
        m2uv[r] = {left_1[i], right_1[i]}

    m2uv[left[-1]] = {7131} #7133

    p = uv_map[7133, :]
    closest = None
    min_dist = 10000
    for i in range(uv_map.shape[0]):
        dist = np.sqrt((p[0] - uv_map[i, 0])**2 + (p[1] - uv_map[i, 1]))
        if dist < min_dist and i != 7133:
            min_dist = dist
            closest = i

    mesh_coord = None
    for key in m2uv.keys():
        if len(list(m2uv[key])) == 2:
            continue
        if list(m2uv[key])[0] == closest:
            mesh_coord = key

    assert m2uv[mesh_coord] == { closest }

    m2uv[mesh_coord] = { closest, 7133 }

    np.save('custom_utils/m2uv_skirt.npy', m2uv)

    # m2uv = m2uv_array(m2uv)
    # uv_counter = np.zeros((uv_map.shape[0]))
    # uv_counter[m2uv[:, 0]] = 1
    # uv_counter[m2uv[:, 1]] += 1

    print()