# check if upscaler is better then filling
import os

from models.upscaler import UpscalerBasic as Upscaler
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

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

from custom_utils.custom_utils import UV2Mesh, _maskFaces, load_model
# 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)

    upscale_enhancer = Upscaler(512, res)
    load_model(upscale_enhancer, 'checkpoints/up_2_3')

    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_low = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape_low = uv_unposed_no_shape_low['uv_unposed_no_shape']
        uv_unposed_no_shape_filled_low = np.load(os.path.join(folder, 'uv_unposed_no_shape_filled' + res_modifier, id + '_' + frame_nr + '.npz'));  uv_unposed_no_shape_filled_low = uv_unposed_no_shape_filled_low['uv_unposed_no_shape_filled']
        mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy'))
        uv_mask = np.load(os.path.join(folder, 'uv_mask' + res_modifier, id + '.npz')); uv_mask = uv_mask['uv_mask']
        mesh = uv2mesh(uv_unposed_no_shape_low[np.newaxis,:,:,:])[0].numpy()
        mesh_filled = uv2mesh(uv_unposed_no_shape_filled_low[np.newaxis,:,:,:])[0].numpy()
        mesh_orig = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_orig = mesh_orig['mesh_unposed_no_shape']

        uv_enhanced = upscale_enhancer([uv_unposed_no_shape_low[np.newaxis,:,:,:], uv_mask[np.newaxis,:,:].astype(np.float32)], training=False)
        mesh_enhanced = uv2mesh(uv_enhanced)[0].numpy()

        metric = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh) ** 2, axis=-1))) / np.sum(mesh_mask)
        metric_filled = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh_filled) ** 2, axis=-1))) / np.sum(mesh_mask)
        metric_enhanced = np.sum(np.sqrt(np.sum(mesh_mask[:,np.newaxis] * (mesh_orig - mesh_enhanced) ** 2, axis=-1))) / np.sum(mesh_mask)
        print(metric)
        print(metric_filled)
        print(metric_enhanced)
        print()

        mesh, faces = _maskFaces(mesh, F, mesh_mask)
        mesh_filled, _ = _maskFaces(mesh_filled, F, mesh_mask)
        mesh_enhanced, _ = _maskFaces(mesh_enhanced, F, mesh_mask)
        x = mesh
        y = mesh_filled
        z = mesh_enhanced
        y[:, 2] += 0.5
        z[:, 2] += 1
        meshes = [go.Mesh3d(
                    x=x[:,0],
                    y=x[:,2],
                    z=x[:,1],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                    ),
                    go.Mesh3d(
                    x=y[:,0],
                    y=y[:,2],
                    z=y[:,1],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                    ),
                    go.Mesh3d(
                    x=z[:,0],
                    y=z[:,2],
                    z=z[:,1],
                    i = faces[:,0],
                    j = faces[:,1],
                    k = faces[:,2]
                    )]

        fig = go.Figure(data=meshes)
        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
        fig.update_layout(
            scene = dict(
                xaxis = dict(nticks=4, range=[-1,1],),
                            yaxis = dict(nticks=4, range=[-1,3],),
                            zaxis = dict(nticks=4, range=[-1,1],),)
                            )

        fig.show()

        print()