import os
import sys
import numpy as np
import trimesh as tm
import scipy.io as sio
import plotly.graph_objects as go
import plotly.express as px
from tqdm import tqdm
import torch
from scipy.spatial.transform import Rotation as R
import argparse

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')

from data.cloth3d_io import readOBJ, readPC2
from data.utils import Frame
from smpl.smpl_torch import SMPL


def calculate_normals(vertices, mesh_file, poses):
    _, faces, _, _ = readOBJ(mesh_file)
    rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)

    normals = []
    for i in range(vertices.shape[0]):
        r = R.from_matrix(rodr[i, 0, :, :])
        # assert r.as_matrix() == rodr[frame, 0, :, :]
        out = r.apply(vertices[i], inverse=True)

        # create a plotly mesh plot of out and vertices [i]
        # fig = go.Figure(data=[
        #     go.Mesh3d(
        #         x=out[:, 0], y=out[:, 2], z=out[:, 1],
        #         # i, j and k give the vertices of triangles
        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
        #         name='y',
        #         showscale=False
        #     ),
        #     go.Mesh3d(
        #         x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
        #         # i, j and k give the vertices of triangles
        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
        #         name='x',
        #         showscale=False
        #     )
        # ])
        # fig.show()
    
        mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
        normals.append(mesh.vertex_normals)
    return np.stack(normals)


def get_frames(folder):
    uv_static_folder = os.path.join(folder, 'pose_shape') 
    frames = sorted(os.listdir(uv_static_folder))

    frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
    frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
    return frames


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--folder", help = "end of range", default='')
    parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/')
    parser.add_argument("-s", "--start", help = "start of range", default=0)
    parser.add_argument("-e", "--end", help = "end of range", default=10)
    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
    args = parser.parse_args()

    obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
    folder = args.folder #'/data-local/127/train'

    os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
    frames = get_frames(folder)
    smpl = SMPL('cpu', 'smpl')

    ids = {frame.id for frame in frames}
    ids = sorted(list(ids))[int(args.start):int(args.end)]

    for id in tqdm(ids):
        # find obj files in the folder
        garment_obj_folder = os.path.join(obj_folder, id[:5])
        garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]

        if len(id) == 5:
            garment = garments[0]
        elif len(id) == 6:
            garment = garments[1]
        else:
            raise ValueError('Invalid frame id')

        mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
        vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
        poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
        normals = calculate_normals(vertices, mesh_file, poses)

        map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))

        mesh_normals = normals[:,map,:]
        mesh_normals = mesh_normals * mask[None,:, None]

        for i in range(vertices.shape[0]):
            # np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
            np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])

        # pass

        # print(frame.id, frame.frame_nr)