import os
import sys
import random
import numpy as np
import tables as tb
import plotly.graph_objects as go
import plotly.express as ex
from plotly.subplots import make_subplots

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame
from custom_utils.custom_utils import _maskFaces


def get_frames(folder, frames_file):
    frames = []
    with open(os.path.join(folder, frames_file), 'r') as f:
        for line in f:
            frames.append(line.strip())
    return frames


def verify_uv(folder, frame):
    id, frame_nr, garment = frame.split()
    uv_static = np.load(os.path.join(folder, 'uv_static', f'{id}_{garment}.npz'))['uv_static']
    uv_mask = np.load(os.path.join(folder, 'uv_mask', f'{id}_{garment}.npz'))['uv_mask']
    uv_posed = np.load(os.path.join(folder, 'uv_posed' + orient_mod, f'{id}_{frame_nr}_{garment}.npz'))['uv_posed']
    uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed' + orient_mod, f'{id}_{frame_nr}_{garment}.npz'))['uv_body_posed']
    uv_curvature = np.load(os.path.join(folder, 'uv_curvature', f'{id}_{frame_nr}_{garment}.npz'))['uv_curvature']

    all_1 = np.concatenate([uv_static, uv_posed, uv_body_posed], axis=1)
    all_1 = (all_1 - all_1.min()) / (all_1.max() - all_1.min())

    ex.imshow(all_1).show()
    ex.imshow(uv_mask).show()
    ex.imshow(uv_curvature).show()


def verify_mesh(folder, frame):
    id, frame_nr, garment = frame.split()
    frame_nr = int(frame_nr)
    info = loadInfo(os.path.join(folder, id, 'info.mat'))
    mask = read_garment_mask(folder, id, garment)
    mesh_static = read_garment_template(folder, id, garment, mask)[mask,:]
    mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=mask, shape=info['shape'], pose=info['poses'][:,frame_nr], deorient=use_no_orient, unshape=no_shape)[mask,:]
    body_posed = read_posed_body_frame(folder, id, frame_nr, info['poses'], use_no_orient)
    mesh_curvature = read_curvature_frame(folder + '_extras', id, garment, frame_nr)[mask]

    fig = make_subplots(rows=1, cols=4, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(4)] ])
    F = np.load(f'custom_utils/smpl_faces{skirt_modifier}.npy').astype(np.int32)
    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)


    fig.add_trace(go.Mesh3d(
                x=mesh_static[:,0],
                y=mesh_static[:,1],
                z=mesh_static[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),
            row=1, col=1)
    
    fig.add_trace(go.Mesh3d(
                x=mesh_posed[:,0],
                y=mesh_posed[:,1],
                z=mesh_posed[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),
            row=1, col=2)

    fig.add_trace(go.Mesh3d(
                x=body_posed[:,0],
                y=body_posed[:,1],
                z=body_posed[:,2],
                i = F[:,0],
                j = F[:,1],
                k = F[:,2],
            ),
            row=1, col=3)

    fig.add_trace(go.Mesh3d(
                x=mesh_posed[:,0],
                y=mesh_posed[:,1],
                z=mesh_posed[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2],
                vertexcolor=np.tile(mesh_curvature[:,None], (1,3)) * 40,
            ),
            row=1, col=4)

    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
                    scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
                    scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
                    scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
                    )
    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
                yaxis = dict(nticks=4, range=[-1,1],),
                zaxis = dict(nticks=4, range=[-1,1],),)
    fig.update_layout(scene = helper,
                        scene2 = helper,
                        scene3 = helper,
                        scene4 = helper,
                        )

    fig.show()

    pass


if __name__ == '__main__':
    folder = '/home/workspace/data/dataset/cloth3d/train'
    frames = get_frames(folder, 'frames_all.txt')
    root = '/'.join(folder.split('/')[:-1])
    type = folder.split('/')[-1]
    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')

    use_no_orient = False
    no_shape = False
    is_skirt = True

    skirt_modifier = '_skirt' if is_skirt else ''
    orient_mod = '_no_orient' if use_no_orient else ''
    shape_mod = '_no_shape' if no_shape else ''


    while True:
        frame = random.choice(frames)
        frame = '05936 0 Dress'

        verify_mesh(folder, frame)
        verify_uv(folder + '_extras', frame)

        pass