import os
import sys
from configs.iccv.config_abl_pyr_32_64_128 import config

os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
# import tensorflow as tf

# gpus = tf.config.experimental.list_physical_devices('GPU')
# if gpus:
#     try:
#         # Currently, memory growth needs to be the same across GPUs
#         for gpu in gpus:
#             tf.config.experimental.set_memory_growth(gpu, True)
#         logical_gpus = tf.config.experimental.list_logical_devices('GPU')
#         print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
#     except RuntimeError as e:
#         # Memory growth must be set before GPUs have been initialized
#         print(e)

import csv
import glob
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

# if config['is_skirt']:
#     from losses import edge_loss_skirt as edge_loss
# else:
#     from losses import edge_loss



# from postprocessing import fix_collision
# from losses import collision_ratio, kl_loss #, concept_loss
# # from losses import collision_penalty as collision_penalty
# from losses import collision_loss as collision_penalty
# from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
# from custom_utils.smpl import SMPL
# from custom_utils.supersmpl_io import writeOBJ




# uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--name', type=str)
    args = vars(parser.parse_args())

    # smpl = SMPL(use_one_hot=config['use_one_hot'])
    skirt_modifier = '_skirt' if config['is_skirt'] else ''

    folder = config['test_dir']

    checkpoint_name = config['name']
    # iteration = '200000'

    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
    garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()


    output_folder = '/home/workspace/local-data/results/iccv/' + checkpoint_name + '-'
    # check if there is a folder that starts with this name using glob
    folders = glob.glob(output_folder + '*')
    assert len(folders) == 1, 'There should be only one checkpoint that starts with this name'
    output_folder = folders[0]
    iteration = output_folder.split('-')[-1]


    # faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
    # ids = os.listdir(faces_folder)
    # ids = [id.split('.')[0] for id in ids]
    # faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }

    masks_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/mesh_mask'
    ids = os.listdir(masks_folder)
    ids = [id.split('.')[0] for id in ids]
    masks = {id : np.load(masks_folder + '/' + id + '.npy') for id in ids }

    frames = os.listdir(output_folder)
    frames = [frame for frame in frames if not frame.endswith('_fixed_coll.npy')]

    if config['remove_global_orientation']:
        gt_name = 'mesh_posed_no_orient'
    else:
        gt_name = 'mesh_posed'

    orig_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/' + gt_name

    errors = dict()
    if config['is_skirt']:
        per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
    else:
        per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }

    for frame in tqdm.tqdm(frames):
        id = frame.split('_')[0]
        frame_nr = frame.split('.')[0].split('_')[1]
        gt = np.load(os.path.join(orig_folder, id + '_' + frame_nr + '.npz'))[gt_name][masks[id], :].astype(np.float32)
        pred = np.load(os.path.join(output_folder, id + '_' + frame_nr + '.npy'))
        metric = np.mean(np.sqrt(np.sum((gt - pred) ** 2, axis=-1)))
        errors[id + '_' + frame_nr] = metric
        per_garment_errors[garment_type_table[id]].append(metric)


    for key in per_garment_errors.keys():
        print(key + ': ' + str(np.mean(per_garment_errors[key])))

    with open('/home/workspace/data/results/errors/iccv/' + checkpoint_name + '-' + iteration + '.csv', mode='w') as f:
        file_writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)

        for key in sorted(errors.keys()):
            file_writer.writerow([key, errors[key]])

        # x = gt
        # y = pred
        # fig = go.Figure(data=[
        #         go.Mesh3d(
        #             x=x[:,0],
        #             y=x[:,1],
        #             z=x[:,2],
        #             i = faces[id][:,0],
        #             j = faces[id][:,1],
        #             k = faces[id][:,2],
        #         ),
        #         go.Mesh3d(
        #             x=y[:,0],
        #             y=y[:,1],
        #             z=y[:,2],
        #             i = faces[id][:,0],
        #             j = faces[id][:,1],
        #             k = faces[id][:,2],
        #         ),
        #     ])

        # fig.update_layout(scene_aspectmode='manual',
        #                 scene_aspectratio=dict(x=1, y=1, z=1))
        # fig.update_layout(
        #     scene = dict(
        #         xaxis = dict(nticks=4, range=[-1,1],),
        #                     yaxis = dict(nticks=4, range=[-1,1],),
        #                     zaxis = dict(nticks=4, range=[-1,1],),)
        #                     )

        # fig.show()

    #     print()

    # print()