# Add the path to the partent directory to use local version of package
import os
os.sys.path.insert(1, os.path.dirname(os.path.abspath('.')))
import fairlyuncertain as fu
from tqdm import tqdm
import pickle
# reload
from importlib import reload
reload(fu)

algorithms = {k: fu.algorithms[k] for k in fu.binary_fairness}

num_runs = 10
    
for dataset in tqdm(fu.binary_datasets):
    filename = f'cached/binary_fairness_{dataset}.pkl'
    if os.path.exists(filename): continue
    results = {}
    for num_run in range(num_runs):
        instance = fu.load_instance(dataset)
        results[num_run] = {'instance' : instance}
        for algo_name, algorithm in algorithms.items():
            results[num_run][algo_name] = algorithm(instance)
    pickle.dump(results, open(filename, 'wb'))

abstention_methods = ['Random', 'Ensemble', 'Selective Ensemble', 'Self-(in)consistency', 'Binomial NLL']
exclude_metrics = []

for abstention_method in abstention_methods:
    dataset = 'ACSIncome'
    print()
    print(f'Running {dataset} with {abstention_method}')
    print()
    filename = f'cached/binary_fairness_{dataset}.pkl'
    results = pickle.load(open(filename, 'rb'))

    #use_metrics = ['Error Rate', 'Statistical Parity', 'Equalized Odds']
    metrics = {metric: fu.binary_metrics[metric] for metric in fu.binary_metrics if metric not in exclude_metrics}
    algorithms_current = {k: fu.algorithms[k] for k in algorithms if k not in abstention_methods and k != abstention_method}
    metric_values = fu.compute_binary_fairness_extended(abstention_method, results, algorithms_current, metrics)

    fu.print_table(metric_values, filename=f'tables/binary_fairness_{dataset}_{abstention_method}.tex', print_md=True)