from Benchmarking.benchmarks.SlovakPresidentalElections import load_data, perform_simulation, plot_results_with_diff
from DataPreparation.Preparation import merge_dfs
from Prediction.Submatrix import Predictor
from Testing.TestingUtils import get_samples_for_batching


def czech():
    data = load_data('data_cz', 'ID')

    pure_data = [percentage_df for (_, _, percentage_df) in data]
    names = [name.split('.')[0] for (name, _, _) in data]

    df = merge_dfs(pure_data, names)
    pairs = get_samples_for_batching(df)
    for (hist, test) in pairs:
        datum = data[names.index(test.columns[0])]
        info = datum[1]
        pr = Predictor(hist, info, len(hist.columns))
        preds = perform_simulation(test, info, pr, 40)
        ypred = [p['prediction'] for p in preds]
        count = [p['current_result'] for p in preds]
        x = [p['count_percentage'] for p in preds]
        plot_results_with_diff(ypred, count, x, test.columns[0])

def main():
    czech()


if __name__ == '__main__':
    main()