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_means_anova import draw_means_anova
from analysis.tukey_hsd95_test import tukey_hsd95_test
from schedulers.the_scheduler_dvfs import FECMWS
from schedulers.the_scheduler import ECMWS


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


def get_settings_dir():
    return op.join(op.dirname(__file__), '../data/scheduler_settings')


scheduler2group_chosen = {
    FECMWS: {
        'alpha1_3': 3
    }
}

scheduler2group_name = {
    FECMWS: {
        'TS': lambda i: rf'$TS^{{\prime}}_{i + 1}$',
        'DP': lambda i: rf'$DP_{i + 1}$'
    }
}

scheduler2group_label = {
    FECMWS: {
        'Conft': r'$Conf^{\prime}_{thresh}$',
        'alpha1_3': r'$(\alpha_1^{\prime},\alpha_2^{\prime},\alpha_3^{\prime})$',
        'TS': '任务排序方法',
        'DP': '截止期划分方法'
    }
}





def analyze_schedulers_calibration(ignored_factors=(), figsize=(8, 8), seed=9527):
    schedulers = [FECMWS,ECMWS]
    random.seed(seed)
    for scheduler in schedulers:
        scheduler_name = scheduler.__name__
        results_file_list = []
        for file in os.listdir(get_results_dir()):
            if file.startswith(scheduler_name) and file.endswith('_cali.pkl'):
                scheduler_name = ''.join(file.split('.')[:-1])
                results_file_list.append(op.join(get_results_dir(), file))
        if not results_file_list:
            raise RuntimeError(f"Results not found for {scheduler_name}")
        param_keys = None
        app_data = dict()
        app = 'Montage'
        app_best_cost = dict()
        for results_file in results_file_list:
            with open(results_file, 'rb') as f:
                results = pickle.load(f)
            for result in results:
                (*params, cost, violation, duration, N, n, rho) = result[0]

                dummy = scheduler(*params)

                if not param_keys:
                    param_keys = dummy.get_params_dict().keys()

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

                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')

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

                app_data[app].append(
                    {
                        **dummy.get_params_dict(),

                        'app': app,
                        'N': N,
                        'n': n,
                        'rho': rho,
                        'duration': duration,
                        'violation': violation,
                        'cost': cost,
                        'rpd': 0,
                    }
                )


        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
        data = [row for app_rows in app_data.values() for row in app_rows]
        df = pd.DataFrame(data)
        df.to_csv(op.join(get_results_dir(), f'{scheduler_name}_rpd.csv'), index=False)
        for factor in param_keys:
            if factor in ignored_factors:
                continue

            factor_df = df.groupby(factor)['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'{scheduler_name}_{factor}_means.csv'))
            fig, ax = plt.subplots(figsize=figsize)
            means, conf_intervals = tukey_hsd95_test(df[[factor, 'rpd']], factor, 'rpd')
            num_groups = df[factor].nunique()
            if scheduler in scheduler2group_chosen and factor in scheduler2group_chosen[scheduler]:
                num_groups = scheduler2group_chosen[scheduler][factor]
            group2name = None
            if scheduler in scheduler2group_name and factor in scheduler2group_name[scheduler]:
                group2name = scheduler2group_name[scheduler][factor]
            label = factor
            if scheduler in scheduler2group_label and factor in scheduler2group_label[scheduler]:
                label = scheduler2group_label[scheduler][factor]
            draw_means_anova(ax, label, means, conf_intervals, num_groups, group2name=group2name)
            fig.savefig(op.join(get_results_dir(), f'{scheduler_name}_{factor}_means.png'), dpi=300)


if __name__ == '__main__':
    analyze_schedulers_calibration()