import os
import sys
import cv2
import numpy as np
import trimesh as tm
import matplotlib.pyplot as plt

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


def uv_to_mesh(uv):
    uv_res = 256
    # m2uv = np.load('/workspace/Code/featurelines/MeshUtils/m2uv.npy', allow_pickle=True).item()
    m2uv = np.load(r'D:\Work\Code\featurelines\MeshUtils\m2uv.npy', allow_pickle=True).item()
    m2uv = m2uv_array(m2uv)
    # uv_map = np.load('/workspace/Code/featurelines/Meshutils/stretched_uv_map_3.npy')
    uv_map = np.load(r'D:\Work\Code\featurelines\MeshUtils\stretched_uv_map_3.npy')
    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)

    mesh = np.zeros((14475, 3))
    uv_pixels = uv_map_discreet[m2uv]
    # counter = np.ones((14475))
    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
    return mesh


def plot(uv, mesh):
    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(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
    # plt.show()

    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(111)
    ax.imshow(np.round( (uv - uv.min()) / (uv.max() - uv.min()) * 255).astype(np.int32))
    plt.show()



if __name__ == '__main__':
    # folder =  '/workspace/Checkpoints/checkpoint_32100/output'
    folder = r'D:\Work\Checkpoints\checkpoint_52000\output'

    files = os.listdir(folder)

    for file in files:
        if not file.endswith('.npz'):
            continue

        uv = np.load(os.path.join(folder, file))
        uv = np.transpose(uv['uv'][0], [1,2,0])
        # uv[abs(uv) < 0.01] = 0
        # print(np.sum(abs(uv) < 0.01))
        mesh = uv_to_mesh(uv)

        plot(uv, mesh)

        pass