import pickle
from collections import defaultdict

from joblib import Parallel, delayed
from tqdm import tqdm
import pandas as pd
from glob import glob
from pathlib import Path
import quapy as qp
from quapy.method.confidence import ConfidenceEllipseSimplex, ConfidenceEllipseCLR, ConfidenceEllipseILR

pd.set_option('display.max_columns', None)
pd.set_option('display.width', 2000)
pd.set_option('display.max_rows', None)
pd.set_option("display.expand_frame_repr", False)
pd.set_option("display.precision", 4)
pd.set_option("display.float_format", "{:.4f}".format)


def compute_coverage_amplitude(region_constructor):
    all_samples = results['samples']
    all_true_prevs = results['true-prevs']

    def process_one(samples, true_prevs):
        ellipse = region_constructor(samples)
        return ellipse.coverage(true_prevs), ellipse.montecarlo_proportion()

    out = Parallel(n_jobs=3)(
        delayed(process_one)(samples, true_prevs)
        for samples, true_prevs in tqdm(
            zip(all_samples, all_true_prevs),
            total=len(all_samples),
            desc='constructing ellipses'
        )
    )

    # unzip results
    coverage, amplitude = zip(*out)
    return list(coverage), list(amplitude)


def update_pickle(report, pickle_path, updated_dict:dict):
    for k,v in updated_dict.items():
        report[k]=v
    pickle.dump(report, open(pickle_path, 'wb'), protocol=pickle.HIGHEST_PROTOCOL)


def update_pickle_with_region(report, file, conf_name, conf_region_class):
    if f'coverage-{conf_name}' not in report:
        cov, amp = compute_coverage_amplitude(conf_region_class)

        update_fields = {
            f'coverage-{conf_name}': cov,
            f'amplitude-{conf_name}': amp,
        }

        update_pickle(report, file, update_fields)

for setup in ['binary', 'multiclass']:
    path = f'./results/{setup}/*.pkl'
    table = defaultdict(list)
    for file in tqdm(glob(path), desc='processing results', total=len(glob(path))):
        file = Path(file)
        dataset, method = file.name.replace('.pkl', '').split('__')
        report = pickle.load(open(file, 'rb'))
        results = report['results']
        n_samples = len(results['ae'])
        table['method'].extend([method.replace('Bayesian','Ba').replace('Bootstrap', 'Bo')] * n_samples)
        table['dataset'].extend([dataset] * n_samples)
        table['ae'].extend(results['ae'])
        table['c-CI'].extend(results['coverage'])
        table['a-CI'].extend(results['amplitude'])

        update_pickle_with_region(report, file, conf_name='CE', conf_region_class=ConfidenceEllipseSimplex)
        update_pickle_with_region(report, file, conf_name='CLR', conf_region_class=ConfidenceEllipseCLR)
        update_pickle_with_region(report, file, conf_name='ILR', conf_region_class=ConfidenceEllipseILR)

        table['c-CE'].extend(report['coverage-CE'])
        table['a-CE'].extend(report['amplitude-CE'])

        table['c-CLR'].extend(report['coverage-CLR'])
        table['a-CLR'].extend(report['amplitude-CLR'])

        table['c-ILR'].extend(report['coverage-ILR'])
        table['a-ILR'].extend(report['amplitude-ILR'])


    df = pd.DataFrame(table)

    n_classes = {}
    tr_size   = {}
    for dataset in df['dataset'].unique():
        fetch_fn = {
            'binary': qp.datasets.fetch_UCIBinaryDataset,
            'multiclass': qp.datasets.fetch_UCIMulticlassDataset
        }[setup]
        data = fetch_fn(dataset)
        n_classes[dataset] = data.n_classes
        tr_size[dataset] = len(data.training)

    # remove datasets with more than max_classes classes
    max_classes = 30
    for data_name, n in n_classes.items():
        if n > max_classes:
            df = df[df["dataset"] != data_name]

    for region in ['CI', 'CE', 'CLR', 'ILR']:
        if setup == 'binary' and region=='ILR':
            continue
        pv = pd.pivot_table(
            df, index='dataset', columns='method', values=['ae', f'c-{region}', f'a-{region}'], margins=True
        )
        pv['n_classes'] = pv.index.map(n_classes).astype('Int64')
        pv['tr_size'] = pv.index.map(tr_size).astype('Int64')
        pv = pv.drop(columns=[col for col in pv.columns if col[-1] == "All"])
        print(f'{setup=}')
        print(pv)
        print('-'*80)