import os
import csv
import glob
import numpy as np


if __name__ == '__main__':
    garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
    folder = '/home/workspace/data/results/errors/iccv'

    checkpoint_name = 'iccv_pyr_32_64_128_skin'
    checkpoint_name_skirt = 'iccv_pyr_32_64_128_skin_skirt'
    iteration = '75' + '000'
    iteration_skirt : str = '155' + '000'
    error_type = 'collision' # 'euclidean' 'collision' edge normal


    errors = []
    errors_skirt = []
    per_garment_errors= { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
    per_garment_errors_skirt  = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }


    # with open(os.path.join(folder, f'{checkpoint_name}-{iteration}-{error_type}.csv'), 'r') as f:
    #     csv_reader = csv.reader(f, delimiter=',')
    #     for row in csv_reader:
    #         per_garment_errors[garment_type_table[row[0].split('_')[0]]].append(float(row[1]))

    # print('\nNon-skirt:')
    # for key in per_garment_errors.keys():
    #     avg = np.mean(per_garment_errors[key])
    #     errors.append(avg)
    #     print(key + ': ' + str(avg))

    # print('\nAvg non-skirt: ' + str(np.mean(errors)))

    with open(os.path.join(folder, f'{checkpoint_name_skirt}-{iteration_skirt}-{error_type}.csv'), 'r') as f:
        csv_reader = csv.reader(f, delimiter=',')
        for row in csv_reader:
            per_garment_errors_skirt[garment_type_table[row[0].split('_')[0]]].append(float(row[1]))

    print('\nSkirt:')
    for key in per_garment_errors_skirt.keys():
        avg = np.mean(per_garment_errors_skirt[key])
        errors_skirt.append(avg)
        print(key + ': ' + str(avg))

    print('Avg skirt: ' + str(np.mean(errors_skirt)))

    # cum_avg_shirt = np.mean(per_garment_errors['Tshirt'] + per_garment_errors_skirt['Tshirt'])
    # cum_avg_top = np.mean(per_garment_errors['Top'] + per_garment_errors_skirt['Top'])

    # print('Global tshirt avg: ' + str(cum_avg_shirt))
    # print('Global top avg: ' + str(cum_avg_top))

    # print('\nGlobal per garment: ')

    # global_errors = []
    # per_garment_errors['Tshirt'].extend(per_garment_errors_skirt['Tshirt'])
    # per_garment_errors['Top'].extend(per_garment_errors_skirt['Top'])
    # per_garment_errors['Skirt'] = per_garment_errors_skirt['Skirt']
    # per_garment_errors['Dress'] = per_garment_errors_skirt['Dress']
    # for key in per_garment_errors.keys():
    #     avg = np.mean(per_garment_errors[key])
    #     global_errors.append(avg)
    #     print(key + ': ' + str(avg), 'samples: ' + str(len(per_garment_errors[key])))

    # print('Global avg: ' + str(np.mean(global_errors)))
    # print(len(global_errors))