import numpy as np
import os
os.sys.path.insert(1, os.path.dirname(os.path.abspath('.')))
import fairlyuncertain as fu
# reload
from importlib import reload
reload(fu)

from sklearn.model_selection import KFold

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

dataset = 'ACSIncome'

k = 5
results = {}
real_dataset = fu.read_dataset(dataset)
n = real_dataset['n'] 
# Ensure n is divisible by k
n = n - n % k
instance = {
    'name' : dataset,
    'n': real_dataset['n'],
    'synthetic': False,
    'X': real_dataset['X'],
    'y': real_dataset['y'],
    'group' : real_dataset['group'],        
    'max_depth' : 6,
    'gamma' : 0,
    'num_boost_round' : 10,
}
for i, (train_idx, test_idx) in enumerate(KFold(n_splits=k).split(range(n))):
    instance['train_indices'] = train_idx
    instance['group_train'] = real_dataset['group'][train_idx]
    instance['group_test'] = real_dataset['group'][test_idx]
    instance['X_train'] = real_dataset['X'][train_idx]
    instance['X_test'] = real_dataset['X'][test_idx]
    instance['y_train'] = real_dataset['y'][train_idx]
    instance['y_test'] = real_dataset['y'][test_idx]
    results[i] = {'instance' : instance}
    for algo_name, algorithm in algorithms.items():
        results[i][algo_name] = algorithm(instance)

metric_values = fu.compute_binary_fairness(results, algorithms, fu.binary_metrics)

fu.print_table(metric_values, filename=f'tables/binary_fairness_cv_{dataset}.tex')