import os
import pickle
import random
from os import path as op

import pandas as pd
from matplotlib import pyplot as plt

from analysis.draw_interaction_anova import draw_interaction_anova
from analysis.tukey_hsd95_test import tukey_hsd95_test_interaction
from baseline.baseline_heft import HEFT
from baseline.baseline_scheduler import DEWS, ECWSD
from schedulers.the_scheduler import ECMWS


def get_results_dir():
    return op.join(op.dirname(__file__), '../results/comparing')


factor2label = {
    'rho': r'$\rho$'
}

scheduler2name = {
    ECMWS: 'ECMWS',
    DEWS: 'DEWS',
    ECWSD: 'ECWSD'
}

group2name = {
    'n': {
        60: 50,
        120: 100,
        250: 250,
        500: 500
    }
}


def analyze_schedulers_comparison(
        ignored_apps=('Montage',),
        ignored_factors=(),
        figsize=(8, 8), seed=9527):
    plt.rcParams.update({"axes.unicode_minus": False})
    best_schedulers = [ECMWS, ECWSD, DEWS, HEFT]
    results_file_list = []
    random.seed(seed)
    for cls in best_schedulers:
        for file in os.listdir(get_results_dir()):
            if file == f'{cls.__name__}_comp.pkl':
                results_file_list.append(op.join(get_results_dir(), file))
                break
    if len(results_file_list) < len(best_schedulers):
        raise RuntimeError(f"Not enough results found for {best_schedulers}")
    app_data = {}
    app_best_cost = dict()
    app2result = dict()
    rf2results = dict()
    for scheduler, results_file in zip(best_schedulers, results_file_list):
        with open(results_file, 'rb') as f:
            results = pickle.load(f)
        rf2results[results_file] = results
        for i, result in enumerate(results):
            app, cost, violation, duration, N, n, rho, name = result

            if app in ignored_apps:
                continue

            if app not in app_best_cost:
                app_best_cost[app] = dict()
            joint_factors = f'{N}_{n}_{rho}'
            if joint_factors not in app_best_cost[app]:
                app_best_cost[app][joint_factors] = float('inf')

            if scheduler in scheduler2name and name == scheduler.__name__:
                name = scheduler2name[scheduler]

            app_best_cost[app][joint_factors] = min(cost, app_best_cost[app][joint_factors])

            if app not in app_data:
                app_data[app] = []
                app2result[app] = []

            app_data[app].append(
                {
                    'app': app,
                    'name': name,
                    'N': N,
                    'n': n,
                    'rho': rho,
                    'duration': duration,
                    'violation': violation,
                    'cost': cost,
                    'rpd': 0,
                }
            )

    joint_schedulers_name = '_'.join([scheduler2name.get(cls, cls.__name__) for cls in best_schedulers])

    for app, data in app_data.items():
        for row in data:
            best_cost = app_best_cost.get(app).get(f"{row['N']}_{row['n']}_{row['rho']}")
            row['rpd'] = (row['cost'] - best_cost) / best_cost * 100

        df = pd.DataFrame(data)
        df.to_csv(op.join(get_results_dir(), f'{joint_schedulers_name}_{app}_rpd.csv'), index=False)

        for factor in ['N', 'n', 'rho']:
            if factor in ignored_factors:
                continue
            factor_df = df.groupby([factor, 'name'])['rpd'].agg(['mean', 'sem'])
            factor_df.columns = ['mean', 'err']
            factor_df = factor_df.reset_index()
            factor_df.to_csv(op.join(get_results_dir(), f'{joint_schedulers_name}_{factor}_{app}_means.csv'))
            fig, ax = plt.subplots(figsize=figsize)
            num_groups = df[factor].nunique()
            mean, ci_lower, ci_upper = tukey_hsd95_test_interaction(df, factor, 'name', 'rpd')
            label = factor2label.get(factor, factor)

            draw_interaction_anova(ax, label, mean, ci_lower, ci_upper, num_groups)
            if factor == 'n':
                factor = 'nn'
            fig.savefig(op.join(get_results_dir(), f'{joint_schedulers_name}_{factor}_{app}_means.png'), dpi=300)


if __name__ == '__main__':
    analyze_schedulers_comparison()