# check if downscaled mask amtches the uv maps, not losing information

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"

import tqdm
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex

from custom_utils.custom_utils import UV2Mesh, _maskFaces
# from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces


if __name__ == '__main__':
    folder = '/home/workspace/local-data/train_100'
    mask = np.load('custom_utils/mask_256x512.npy')
    F = np.load("custom_utils/smpl_faces.npy")
    uv2mesh = UV2Mesh(res=256)
    res = 32
    res_modifier = '_' + str(res)


    frames = os.listdir(os.path.join(folder, 'pose_shape'))
    visited = set()
    meshes = []
    count = 0

    for frame in tqdm.tqdm(frames):
        id, frame_nr = frame.split('_')
        frame_nr = frame_nr.split('.')[0]
        if id in visited:
            continue
        visited.add(id)
        print(id)
        uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
        uv_mask = np.load(os.path.join(folder, 'uv_mask' + res_modifier, id + '.npz')); uv_mask = uv_mask['uv_mask']

        check_1 = uv_unposed_no_shape * uv_mask[:,:,np.newaxis]
        check_2 = uv_unposed_no_shape * uv_mask[:,:,np.newaxis] - uv_unposed_no_shape

        img = np.concatenate((uv_unposed_no_shape, np.tile(uv_mask[:,:,np.newaxis], [1,1,3]), check_1, check_2), axis=1)

        ex.imshow(img).show()

        print()



# WORKS!