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


os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = '0'


import time
import argparse
import wandb
import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image


import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm import tqdm
from data.dataset import Dataset
from data.utils import Frame
from data.cloth3d_io import readPC2Frame, readPC2
from data.preprocessing import deorient
from utils.checkpoints import CheckpointSaver, load_model
from torch.utils.data import DataLoader
from utils.utils import text_to_html, pickle_load
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss, kl_loss
from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels
from plotly.subplots import make_subplots
import pytorch3d.loss
from scipy.spatial.distance import cdist
from scipy.spatial import cKDTree
from smpl.smpl_torch import SMPL
from scipy.spatial.transform import Rotation as R


def chamfer(points1, points2, points2_tree: cKDTree):
    tree1 = cKDTree(points1)
    dist1, _ = tree1.query(points2, workers=40)
    dist2, _ = points2_tree.query(points1, workers=40)
    chamfer_dist = np.mean(dist1) + np.mean(dist2)
    return chamfer_dist


smpl = SMPL(path='smpl', device='cpu')
def deorient(verts, pose, inverse=True):
    assert pose.shape[1] == 72
    rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
    rotated = np.zeros_like(verts)
    for frame in range(verts.shape[0]):
        r = R.from_matrix(rodr[frame, 0, :, :])
        rotated[frame] = r.apply(verts[frame], inverse=inverse)

    return rotated

def calculate():
    errors_hood = []
    errors_ours = []
    garments = sorted(os.listdir(hood_folder))
    garments = [g[:-4] for g in garments]

    for garment in tqdm(garments):
        picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
        id, garment_type = garment.split('_')

        pred_hood = pickle_load(f'{hood_folder}/{garment}.pkl')['verts']
        gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
        mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
        pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
        gt = deorient(gt, pose, True)


        for i in range(pred_hood.shape[0]):
            gt_tree = cKDTree(gt[i])

            try:
                pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
            except:
                print('missed one')
                continue

            errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree)))
            errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask,:], gt[i], gt_tree)))

            # fig = go.Figure()
            # fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
            # fig.add_trace(go.Scatter3d(x=pred_hood[i][:, 0], y=pred_hood[i][:, 2], z=pred_hood[i][:, 1], mode='markers', marker=dict(size=1)))
            # fig.update_layout(scene_aspectmode='cube')
            # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
            # fig.show()

            # fig = go.Figure()
            # fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
            # fig.add_trace(go.Scatter3d(x=pred_ours[i,mask,:][:, 0], y=pred_ours[i,mask,:][:, 2], z=pred_ours[i,mask,:][:, 1], mode='markers', marker=dict(size=1)))
            # fig.update_layout(scene_aspectmode='cube')
            # fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
            # fig.show()


    # save errors to a csv file
    with open(f'{results_folder}/hood/{name}_hood.csv', 'w') as f:
        for frame, error in errors_hood:
            frame_str = ','.join(str(frame).split('_'))
            f.write(f'{frame_str},{error}\n')

    with open(f'{results_folder}/hood/{name}_ours.csv', 'w') as f:
        for frame, error in errors_ours:
            frame_str = ','.join(str(frame).split('_'))
            f.write(f'{frame_str},{error}\n')



if __name__ == '__main__':

    name = 'hood_2'

    hood_folder = '/data-local/127_2/results/hood/all_2'
    ours_folder = dict()
    ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
    ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
    cloth3d_folder = '/data-net/dataset/cloth3d/test'

    results_folder = '/data-local/127_2/results'
    out_folder = f'{results_folder}/errors/{name}'


    calculate()