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

import time
import argparse
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from data.utils import Frame




def calculate(type='chamfer'):
    errors = dict()
    errors_mean = dict()

    files_hood = sorted(os.listdir(f'{folder}/{run}_hood/{type}'))
    errors_hood = [np.load(f'{folder}/{run}_hood/{type}/{file}') for file in tqdm(files_hood)]
    files_ours = sorted(os.listdir(f'{folder}/hood_3_var_ours/{type}'))
    errors_ours = [np.load(f'{folder}/hood_3_var_ours/{type}/{file}') for file in tqdm(files_ours)]

    for error, frame in zip(errors_hood, files_hood):
        frame_id, frame_nr, frame_garment = frame.split('_')
        category = frame_garment

        if category not in errors:
            errors[category] = []
        errors[category].append(error)

    for category in errors:
        errors_mean[category] = round(np.mean(errors[category]), 5)
    print(errors_mean) 

    mean = round(np.mean(list(errors_mean.values())), 5)
    print(f'Final HOOD error_{type}: {mean}')

    # errors = dict()
    # errors_mean = dict()

    # for error, frame in zip(errors_ours, files_ours):
    #     frame_id, frame_nr, frame_garment = frame.split('_')
    #     category = frame_garment

    #     if category not in errors:
    #         errors[category] = []
    #     errors[category].append(error)

    # for category in errors:
    #     errors_mean[category] = round(np.mean(errors[category]), 5)
    # print(errors_mean) 

    # mean = round(np.mean(list(errors_mean.values())), 5)
    # print(f'Final OURS error_{type}: {mean}')


if __name__ == '__main__':

    folder = '/data-local/127_2/results/hood'
    run = 'hood_3'

    for type in ['chamfer', 'area', 'collision_ratio']: # ['chamfer', 'collision', 'area', 'collision_ratio']:
        calculate(type)