import os
import re

import numpy as np
import pandas as pd
import argparse

from experimentlogger.experiment_logger import ExperimentLogger

import matplotlib
if matplotlib.get_backend() == 'TkAgg':
    matplotlib.use('Agg')
from shsplots.matplotlib.plot_context import PlotContext

from prolothar_common.experiments.result_collector import ResultCollector

class CreatePlotsSingleQueueExperiment:

    def __init__(self, experiment_folders: list[str], cuemin_analysis: bool = False):
        self.logger = ExperimentLogger(
            'experiments/evaluation/results/reconstruction_single_queue', 'plots')
        self.__experiment_folders = experiment_folders
        self.__cuemin_analysis = cuemin_analysis

    def run(self):
        result_folder_list = []
        for i, experiment_folder in enumerate(self.__experiment_folders):
            result_folder = self.logger.get_absolute_path(str(i))
            result_folder_list.append(result_folder)
            os.mkdir(result_folder)
            self.logger.log(f'plot results of "{experiment_folder}" ({i})')
            self.__evaluate_results_per_experiment(experiment_folder, result_folder)
            if self.__cuemin_analysis:
                self.__analyze_cuemin_log(experiment_folder, result_folder)
        self.__evaluate_results_across_experiments(experiment_folder, result_folder_list)
        self.logger.close()

    def __evaluate_results_across_experiments(self, experiment_folder: str, result_folder_list: list[str]):
        df = pd.concat([
            pd.read_csv(os.path.join(result_folder, 'results.csv'), sep=',')
            for result_folder in result_folder_list
        ])
        if 'timesteps' in experiment_folder:
            group_by_attribute = 'train_timesteps'
        elif 'nr_of_jobs_per_server' in experiment_folder:
            group_by_attribute = 'nr_of_processed_jobs_per_server'
        else:
            raise NotImplementedError()

        result_collector = ResultCollector(group_by_attribute)
        self.__configure_result_collector(result_collector)

        for _, row in df.iterrows():
            for column in ['MAE_c', 'MAPE_c', 'MAE_mean_batch_size', 'accuracy_D', 'inference_time_in_s',
                           'nr_of_models', 'mean_c', 'max_c', 'MAE_mean_waiting_time', 'MAE_mean_sojourn_time',
                           'mean_waiting_time_predicted', 'mean_sojourn_time_predicted',
                           'MAE_departure_time_train', 'MAE_departure_time_test',
                           'MAPE_departure_time_train', 'MAPE_departure_time_test']:
                if self.__should_skip_value_evaluation_for_model(row['model'], column):
                    continue
                match (row[column], row['model']):
                    # a list of numbers, e.g. "[32, 42, 0, 1, 3]"
                    case (str() as value, _ as model_name) if value.startswith('['):
                        for parsed_value in filter(np.isfinite, map(float, value[1:-1].split(', '))):
                            result_collector.add_result(model_name, column, row[group_by_attribute], parsed_value)
                    # a single value, e.g. "42.3"
                    case (float()|int() as value, _ as model_name) if not np.isnan(value):
                        result_collector.add_result(model_name, column, row[group_by_attribute], value)

        all_directory = self.logger.get_absolute_path('all')
        os.mkdir(all_directory)
        result_collector.save_plots(all_directory)

    def __should_skip_value_evaluation_for_model(self, model_name: str, column: str) -> bool:
        return (
            model_name == 'PSM-Batch-Miner'
            and column not in ['MAE_mean_batch_size', 'inference_time_in_s', 'nr_of_models']
        ) or (
            model_name in ['COrder', 'Deltamax']
            and column not in ['MAE_c', 'bias_c', 'MAPE_c', 'inference_time_in_s', 'nr_of_models']
        ) or (
            model_name in ['ground_truth', 'ground_truth re-run']
            and column in ['MAE_c', 'bias_c', 'MAPE_c', 'MAE_mean_batch_size', 'accuracy_D', 'inference_time_in_s']
        )

    def __evaluate_results_per_experiment(self, experiment_folder: str, result_folder: str):
        df = pd.read_csv(os.path.join(experiment_folder, 'evaluation.csv'), sep=';')
        ground_truth_df = df[df['model'] == 'ground_truth'].copy()
        joint_df = df.merge(ground_truth_df, on='i', suffixes=('_predicted', '_actual'))
        joint_df = joint_df.drop([
            'total_timesteps_predicted',
            'train_timesteps_predicted',
            'model_actual',
            'inference_time_in_s_actual',
            'MAE_departure_time_train_actual',
            'MAE_departure_time_test_actual',
            'MAPE_departure_time_train_actual',
            'MAPE_departure_time_test_actual'], axis=1)
        joint_df = joint_df.rename(columns={
            'total_timesteps_actual': 'total_timesteps',
            'train_timesteps_actual': 'train_timesteps',
            'model_predicted': 'model',
            'inference_time_in_s_predicted': 'inference_time_in_s',
            'MAE_departure_time_train_predicted' : 'MAE_departure_time_train',
            'MAE_departure_time_test_predicted' : 'MAE_departure_time_test',
            'MAPE_departure_time_train_predicted' : 'MAPE_departure_time_train',
            'MAPE_departure_time_test_predicted' : 'MAPE_departure_time_test'
        })
        joint_df['MAE_c'] = (
            joint_df['c_actual'] -
            joint_df['c_predicted']).abs()
        joint_df['bias_c'] = (
            joint_df['c_actual'] -
            joint_df['c_predicted'])
        joint_df['MAPE_c'] = (
            (joint_df['c_actual'] - joint_df['c_predicted']) / joint_df['c_actual']).abs()
        joint_df['MAE_mean_batch_size'] = (
            joint_df['mean_service_batch_size_actual'] -
            joint_df['mean_service_batch_size_predicted']).abs()
        joint_df['MAE_mean_waiting_time'] = (
            joint_df['mean_waiting_time_actual'] -
            joint_df['mean_waiting_time_predicted']).abs()
        joint_df['MAE_mean_sojourn_time'] = (
            joint_df['mean_sojourn_time_actual'] -
            joint_df['mean_sojourn_time_predicted']).abs()
        joint_df['accuracy_D'] = (
            joint_df['D_actual'] ==
            joint_df['D_predicted']).astype(int)
        joint_df['mean_c'] = joint_df['c_predicted']
        joint_df['max_c'] = joint_df['c_predicted']
        if 'nr_of_jobs_per_server' in experiment_folder:
            joint_df['nr_of_processed_jobs_per_server'] = [
                int(s) for s in re.findall(r' \d+ ', experiment_folder)][0]
            joint_group_by = ['model', 'nr_of_processed_jobs_per_server']
        else:
            joint_df['nr_of_processed_jobs_per_server'] = joint_df['nr_of_processed_jobs_actual'] / joint_df['c_actual']
            joint_group_by = ['model', 'total_timesteps', 'train_timesteps']
        joint_df['nr_of_models'] = 1
        result_df = joint_df.groupby(by=joint_group_by).agg({
            'nr_of_models': 'sum',
            'mean_c': 'mean',
            'max_c': 'max',
            'MAE_c': list,
            'bias_c': list,
            'MAPE_c': list,
            'MAE_mean_batch_size': list,
            'MAE_mean_waiting_time': list,
            'MAE_mean_sojourn_time': list,
            'mean_sojourn_time_predicted': list,
            'mean_waiting_time_predicted': list,
            'MAE_departure_time_train': list,
            'MAE_departure_time_test': list,
            'MAPE_departure_time_train': list,
            'MAPE_departure_time_test': list,
            'accuracy_D': list,
            'inference_time_in_s': list
        })
        result_df.to_csv(os.path.join(result_folder, 'results.csv'))
        result_df.to_excel(os.path.join(result_folder, 'results.xlsx'))

        for group_by_attribute in ['c_actual', 'nr_of_processed_jobs_actual', 'nr_of_processed_jobs_per_server']:
            if (group_by_attribute == 'c_actual' and 'timesteps' in experiment_folder) \
            or (group_by_attribute == 'nr_of_processed_jobs_per_server' and 'nr_of_jobs_per_server' in experiment_folder):
                result_collector_mode = ResultCollector.modes.LINE
            else:
                result_collector_mode = ResultCollector.modes.SCATTER
            result_collector = ResultCollector(group_by_attribute, mode=result_collector_mode)
            self.__configure_result_collector(result_collector)

            for _, row in joint_df.groupby(by=['model', group_by_attribute]).agg({
                'accuracy_D': list,
                'MAE_mean_batch_size': list,
                'c_actual': list,
                'MAE_c': list,
                'bias_c': list,
                'MAPE_c': list,
                'MAE_departure_time_train': list,
                'MAE_departure_time_test': list,
                'MAPE_departure_time_train': list,
                'MAPE_departure_time_test': list,
                'nr_of_models': 'sum'
            }).iterrows():
                model_name, nr_of_servers = row.name
                for metric, unparsed_value in row.items():
                    if not self.__should_skip_value_evaluation_for_model(model_name, metric):
                        match unparsed_value:
                            case list() as value_list:
                                for value in value_list:
                                    if not np.isnan(value):
                                        result_collector.add_result(model_name, metric, nr_of_servers, value)
                            case float()|int() as value:
                                if not np.isnan(value):
                                    result_collector.add_result(model_name, metric, nr_of_servers, value)
                            case _ as x:
                                raise NotImplementedError((x, type(x)))

            sub_result_folder = os.path.join(result_folder, group_by_attribute)
            os.mkdir(sub_result_folder)
            result_collector.save_plots(sub_result_folder)

    def __analyze_cuemin_log(self, experiment_folder: str, result_folder: str):
        cuemin_folder = os.path.join(experiment_folder, 'CueMin')
        if os.path.exists(cuemin_folder):
            result_folder = os.path.join(result_folder, 'CueMin')
            os.mkdir(result_folder)
            self.logger.log(f'Start analysis on CueMin directory "{cuemin_folder}"')
            self.__plot_mdl_for_varying_c(cuemin_folder, result_folder)
            self.__plot_mdl_for_varying_c_for_distinct_s(cuemin_folder, result_folder)
        else:
            self.logger.log(f'CueMin directory "{cuemin_folder}" does not exist')

    def __plot_mdl_for_varying_c(self, cuemin_folder: str, result_folder: str):
        result_folder = os.path.join(result_folder, 'mdl_for_varying_c')
        os.mkdir(result_folder)
        for file in [f for f in os.listdir(cuemin_folder) if f.endswith('.csv')]:
            model_index = file.replace('cuemin_', '').replace('.csv', '')
            df = pd.read_csv(os.path.join(cuemin_folder, file), sep=';')
            df = df.groupby(by=['D', 'B', 'c']).agg({'mdl_score': min}).reset_index()
            df = df.groupby(by=['D', 'B']).agg({'c': list, 'mdl_score': list}).reset_index()
            for _, row in df.iterrows():
                file_safe_group_name = "".join(c if c.isalnum() else '-' for c in row['D'][:10] + ' ' + row['B'][:7] + row['B'][-5:])
                plot_path = os.path.join(result_folder, f'{model_index}_{file_safe_group_name}.png')
                with PlotContext(show=False, filepath=plot_path) as plot_context:
                    plot_context.get_axes().plot(row['c'], row['mdl_score'])
                    plot_context.get_axes().set_xlabel('c')
                    plot_context.get_axes().set_ylabel('L')

    def __plot_mdl_for_varying_c_for_distinct_s(self, cuemin_folder: str, result_folder: str):
        result_folder = os.path.join(result_folder, 'mdl_for_varying_c_distinct_S')
        os.mkdir(result_folder)
        for file in [f for f in os.listdir(cuemin_folder) if f.endswith('.csv')]:
            model_index = file.replace('cuemin_', '').replace('.csv', '')
            df = pd.read_csv(os.path.join(cuemin_folder, file), sep=';')
            def map_service_time_to_type(s):
                for t in ['Degenerate', 'Poisson', 'Geometric', 'NegativeBinomial']:
                    if t in s:
                        return t
                return 'Unknown'
            df['S'] = df['S'].map(map_service_time_to_type)
            df = df.groupby(by=['D', 'B', 'S', 'c']).agg({'mdl_score': min}).reset_index()
            df = df.groupby(by=['D', 'B', 'S']).agg({'c': list, 'mdl_score': list}).reset_index()
            for _, row in df.iterrows():
                file_safe_group_name = "".join(
                    c if c.isalnum() else '-' for c in (
                        row['D'][:10] + ' ' +
                        row['B'][:7] + row['B'][-5:] + ' ' +
                        row['S']
                    )
                )
                plot_path = os.path.join(result_folder, f'{model_index}_{file_safe_group_name}.png')
                with PlotContext(show=False, filepath=plot_path) as plot_context:
                    plot_context.get_axes().plot(row['c'], row['mdl_score'])
                    plot_context.get_axes().set_xlabel('c')
                    plot_context.get_axes().set_ylabel('L')

    def __configure_result_collector(self, result_collector: ResultCollector):
        result_collector.set_color('ground_truth', ResultCollector.colors.GOLD)
        result_collector.set_marker('ground_truth', ResultCollector.markers.STAR)
        result_collector.set_color('ground_truth re-run', ResultCollector.colors.GREY)
        result_collector.set_marker('ground_truth re-run', ResultCollector.markers.DIAMOND)
        result_collector.set_color('CueMin', ResultCollector.colors.BLUISH_PURPLE)
        result_collector.set_marker('CueMin', ResultCollector.markers.PLUS)
        result_collector.set_color('DInformedCueMin', ResultCollector.colors.LIGHT_BLUE)
        result_collector.set_marker('DInformedCueMin', ResultCollector.markers.PENTAGON)
        result_collector.set_color('SA-CueMin', ResultCollector.colors.LIGHT_BROWN)
        result_collector.set_marker('SA-CueMin', ResultCollector.markers.PENTAGON)
        result_collector.set_color('PSM-Batch-Miner', ResultCollector.colors.LIGHT_BROWN)
        result_collector.set_marker('PSM-Batch-Miner', ResultCollector.markers.TRIANGLE_DOWN)
        result_collector.set_color('Naive-Lifo-vs-Fifo', ResultCollector.colors.ORANGE)
        result_collector.set_marker('Naive-Lifo-vs-Fifo', ResultCollector.markers.X)
        result_collector.set_color('COrder', ResultCollector.colors.RED)
        result_collector.set_marker('COrder', ResultCollector.markers.TRIANGLE_RIGHT)
        result_collector.set_color('Deltamax', ResultCollector.colors.BLUE)
        result_collector.set_marker('Deltamax', ResultCollector.markers.SQUARE)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        'plots result of queue mining experiments')

    parser.add_argument(
            '--experiments', type=str, nargs='+',
            help='a list of paths to folders of experiment results',
            required=True)

    parser.add_argument(
        '--cuemin_analysis', action='store_true',
        help=('add this flag to enable intensive analysis on CueMin results'),
        required=False, default=False)

    arguments = parser.parse_args()

    if arguments.experiments[0].startswith('"'):
        arguments.experiments = [e.replace('"', '') for e in arguments.experiments]
        arguments.experiments = [e.strip() + ')' for e in ' '.join(arguments.experiments).split(')') if e]
    CreatePlotsSingleQueueExperiment(
        arguments.experiments,
        cuemin_analysis=arguments.cuemin_analysis
    ).run()