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

import importlib
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--iteration', type=str, default='')
parser.add_argument('-g', '--gpu', type=str)
args = vars(parser.parse_args())

if args['debug']:
    from configs.iccvw.pyr_64_128_cape_f_2 import config_1
else:
    config_module = importlib.import_module(args['config'])
    config_1 = config_module.config


import time
import argparse
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():
    errors = dict()
    errors_mean = dict()

    with open(f'{results_folder}/errors/{name}.csv', 'r') as f:
        lines = f.readlines()

    for line in lines:
        line = line.split(',')
        category = line[2]
        if category not in errors:
            errors[category] = []
        errors[category].append(float(line[3]))

    for category in errors:
        errors_mean[category] = np.mean(errors[category])

    print(errors_mean) 

    mean = np.mean(list(errors_mean.values()))

    print(f'Final error: {mean}')



if __name__ == '__main__':
    skirt_modifier = '_skirt' if config_1['is_skirt'] else ''

    name = config_1['name']
    if not config_1['is_cape']:
        frames_file = config_1['test_frames']
        with open(os.path.join(config_1['test_dir'], frames_file), 'r') as f:
            frames = f.readlines()
        frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames]
    else:
        from data.dataset_cape import read_frames
        frames = read_frames(config_1['test_dir'], shuffle=False, is_cape=True)

    results_folder = '/data-local/127_2/results'
    if args['iteration'] == '':
        folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_1['name'])]
        assert len(folders) == 1
        name = folders[0]
    else:
        name = f'{name}_{args["iteration"]}'
    out_folder = f'{results_folder}/{name}'


    calculate()