from collections import defaultdict
from math import factorial
from typing import Callable, Union
import gc
import os
from itertools import product, chain
from random import Random
import csv
import argparse
import json
import numpy as np
import pandas as pd

from prolothar_queue_mining.inference.queue.waiting_area.waiting_area_estimator import WaitingAreaEstimator
from prolothar_queue_mining.prediction.departure_time.departure_time_predictor import DepartureTimePredictor
np.seterr(all='raise')
from prolothar_common.experiments.statistics import Statistics
from sklearn.linear_model import LogisticRegression
import matplotlib
#not sure why this is necessary, but otherwise I sometimes observed a "Fail to allocate bitmap" exception
if matplotlib.get_backend() == 'TkAgg':
    matplotlib.use('Agg')

from statsmodels.tools.sm_exceptions import HessianInversionWarning
import warnings
warnings.filterwarnings("ignore", category=HessianInversionWarning)

from shsplots.matplotlib.plot_context import PlotContext
from shsplots.matplotlib.histogram import plot_histogram

from experimentlogger.experiment_logger import ExperimentLogger
from dhparallel.engine.abstract.computation_engine import ComputationEngine
from dhparallel.engine.single_thread.single_thread import SingleThreadComputationEngine
from dhparallel.engine.multiprocess.multiprocess import MultiprocessComputationEngine

from prolothar_common.experiments.stopwatch import Stopwatch

from prolothar_queue_mining.model.queue import Queue
from prolothar_queue_mining.model.job import Job
from prolothar_queue_mining.model.job.job_to_vector_transformer import JobToVectorTransformer
from prolothar_queue_mining.model.job.job_to_vector_transformer import OneHotEncoder
from prolothar_queue_mining.model.population import InfinitePopulation
from prolothar_queue_mining.model.arrival_process import ArrivalWithDistribution
from prolothar_queue_mining.model.arrival_process import LimitedArrival
from prolothar_queue_mining.model.arrival_process import RecordingArrival
from prolothar_queue_mining.model.arrival_process import BatchArrivalAdapter
from prolothar_queue_mining.model.distribution import Distribution
from prolothar_queue_mining.model.distribution import DiscreteDegenerateDistribution
from prolothar_queue_mining.model.server import Server
from prolothar_queue_mining.model.service_time import ServiceTimeWithDistribution
from prolothar_queue_mining.model.waiting_area import WaitingArea
from prolothar_queue_mining.model.waiting_area import FastLastComeFirstServeWaitingArea as LastComeFirstServeWaitingArea
from prolothar_queue_mining.model.waiting_area import FastFirstComeFirstServeWaitingArea as FirstComeFirstServeWaitingArea
from prolothar_queue_mining.model.waiting_area import PriorityClassWaitingArea
from prolothar_queue_mining.model.waiting_area import PairwisePriorityClassifierWaitingArea
from prolothar_queue_mining.model.waiting_area.pairwise_priority_classifier import SklearnPairwisePriorityClassifier
from prolothar_queue_mining.model.observer.waiting_time import WaitingTimeRecordingObserver
from prolothar_queue_mining.model.observer.sojourn_time import SojournTimeRecordingObserver
from prolothar_queue_mining.model.observer.queue_length import QueueLengthRecordingObserver
from prolothar_queue_mining.model.environment import Environment
from prolothar_queue_mining.model.exit import ListCollectorExit
from prolothar_queue_mining.model.event import Event

from prolothar_queue_mining.inference.queue import QueueMiner
from prolothar_queue_mining.inference.queue import CueMin
from prolothar_queue_mining.inference.queue.cuemin.mdl_score import compute_mdl
from prolothar_queue_mining.inference.queue.nr_of_servers import NrOfServersEstimator
from prolothar_queue_mining.inference.queue.nr_of_servers import Deltamax
from prolothar_queue_mining.inference.queue.nr_of_servers import COrder
from prolothar_queue_mining.inference.queue.nr_of_servers import COrderLcfs
from prolothar_queue_mining.inference.queue.batch.psm_batch_miner import PerformanceSpectrumBatchMiner
from prolothar_queue_mining.inference.queue.waiting_area import NaiveLifoOrFifoWaitingAreaEstimator

from prolothar_queue_mining.prediction.departure_time import QueueDepartureTimePredictor

from experiments.utils import parse_distribution
from experiments.utils import parse_population

class AddServerEvent(Event):

    def __init__(self, time: int, queue: Queue):
        super().__init__(time)
        self.__queue = queue

    def execute(self, environment: Environment):
        new_server = self.__queue.get_servers()[-1].copy()
        new_server.set_current_job(None)
        self.__queue.add_server(new_server)

class RemoveServerEvent(Event):

    def __init__(self, time: int, queue: Queue, min_nr_of_servers: int = 1):
        super().__init__(time)
        self.__queue = queue
        self.__min_nr_of_servers = min_nr_of_servers

    def execute(self, environment: Environment):
        if self.__queue.get_nr_of_servers() > self.__min_nr_of_servers:
            self.__queue.remove_server(-1)

def run_evaluation_on_ith_ground_truth_model(
        parameters: dict, index_and_ground_truth_model: tuple[int, Queue]) -> list[Union[dict, str]]:
    i, ground_truth_model = index_and_ground_truth_model
    nr_timesteps_for_simulation = parameters['nr_timesteps_for_simulation']
    nr_of_jobs_per_server = parameters['nr_of_jobs_per_server']
    test_observation_time_proportion = parameters['test_observation_time_proportion']
    data_log_directory = parameters['data_log_directory']
    experiment_directory = parameters['experiment_directory']
    numerical_attributes = parameters['numerical_attributes']
    categorical_attributes = parameters['categorical_attributes']
    create_plots = parameters['create_plots']
    nr_of_changes_distribution_train = parameters['nr_of_changes_distribution_train']
    nr_of_changes_distribution_test = parameters['nr_of_changes_distribution_test']
    probability_add_server = parameters['probability_add_server']
    probability_remove_server = parameters['probability_remove_server']
    record_cuemin_candidates = parameters['record_cuemin_candidates']
    cuemin_waiting_areas = parameters['cuemin_waiting_areas']
    seed = parameters['seed']

    result = []
    # we use copy because ground truth models can share the same reference
    # to arrival processes etc, which would result in a wrong state for
    # simulations that are started later
    ground_truth_model = ground_truth_model.copy()
    initial_state_ground_truth = ground_truth_model.copy()
    nr_of_simulated_timesteps = run_ground_truth_model(
        ground_truth_model, nr_timesteps_for_simulation, nr_of_jobs_per_server,
        test_observation_time_proportion,
        nr_of_changes_distribution_train=nr_of_changes_distribution_train,
        nr_of_changes_distribution_test=nr_of_changes_distribution_test,
        probability_add_server=probability_add_server,
        probability_remove_server=probability_remove_server,
        random=Random(seed + i))
    if nr_of_jobs_per_server is not None:
        test_observation_time = int(nr_of_simulated_timesteps * test_observation_time_proportion)
        nr_timesteps_for_simulation = nr_of_simulated_timesteps
    else:
        test_observation_time = int(nr_of_simulated_timesteps * test_observation_time_proportion)
    all_observed_arrivals, all_observed_departures, train_arrivals, train_departures, test_arrivals, test_departures = \
        create_data_from_executed_ground_truth_model(ground_truth_model, test_observation_time)
    if parameters['log_data']:
        store_data_to_data_log_directory(all_observed_arrivals, all_observed_departures, i, data_log_directory)
    #we remove all unidentifiable models
    #if no jobs have left the queue, inference does not make sense
    #if not all servers have been used at least once, the ground-truth number of
    #servers cannot be inferred
    if not train_departures or not test_departures \
    or ground_truth_model.get_queue_length_observer().get_max_queue_length() == 0:
        return result

    initial_state_ground_truth.set_arrival_process(
        RecordingArrival(ground_truth_model.get_arrival_process().to_fixed_arrival_process()))
    run_model(initial_state_ground_truth, nr_timesteps_for_simulation, test_observation_time)
    result.append(create_model_evaluation_dict(
        initial_state_ground_truth, 'ground_truth', i,
        nr_timesteps_for_simulation, test_observation_time,
        train_arrivals, test_arrivals,
        train_departures, test_departures,
        dict(all_observed_departures)))
    if create_plots:
        plot_model_execution(initial_state_ground_truth, 'ground_truth', i, experiment_directory)

    evaluate_cuemin(
        i, ground_truth_model, nr_timesteps_for_simulation, test_observation_time,
        experiment_directory, numerical_attributes, categorical_attributes, seed, result,
        all_observed_arrivals, all_observed_departures, train_arrivals, test_arrivals,
        train_departures, test_departures, create_plots, record_cuemin_candidates,
        cuemin_waiting_areas)

    result.append(evaluate_batch_miner(
        train_arrivals, train_departures,
        PerformanceSpectrumBatchMiner(), 'PSM-Batch-Miner',
        i, nr_timesteps_for_simulation, test_observation_time
    ))

    result.append(evaluate_waiting_area_estimator(
        train_arrivals, train_departures,
        NaiveLifoOrFifoWaitingAreaEstimator(), 'Naive-Lifo-vs-Fifo',
        i, nr_timesteps_for_simulation, test_observation_time
    ))

    if 'FCFS' == ground_truth_model.get_waiting_area().get_discipline_name():
        result.append(evaluate_nr_of_servers_estimator(
            train_arrivals, train_departures,
            COrder(), 'COrder',
            i, nr_timesteps_for_simulation, test_observation_time
        ))
        result.append(evaluate_nr_of_servers_estimator(
            train_arrivals, train_departures,
            Deltamax(), 'Deltamax',
            i, nr_timesteps_for_simulation, test_observation_time
        ))
    elif 'LCFS' == ground_truth_model.get_waiting_area().get_discipline_name():
        result.append(evaluate_nr_of_servers_estimator(
            train_arrivals, train_departures,
            COrderLcfs(), 'COrder',
            i, nr_timesteps_for_simulation, test_observation_time
        ))
    elif 'LCFS' in ground_truth_model.get_waiting_area().get_discipline_name():
        result.append(evaluate_nr_of_servers_estimator(
            train_arrivals, train_departures,
            COrderLcfs(), 'COrderNoFit',
            i, nr_timesteps_for_simulation, test_observation_time
        ))
    elif 'FCFS' in ground_truth_model.get_waiting_area().get_discipline_name():
        result.append(evaluate_nr_of_servers_estimator(
            train_arrivals, train_departures,
            COrder(), 'COrderNoFit',
            i, nr_timesteps_for_simulation, test_observation_time
        ))
        result.append(evaluate_nr_of_servers_estimator(
            train_arrivals, train_departures,
            Deltamax(), 'DeltamaxNoFit',
            i, nr_timesteps_for_simulation, test_observation_time
        ))
    gc.collect()
    return result

def evaluate_cuemin(
        i, ground_truth_model, nr_timesteps_for_simulation, test_observation_time,
        experiment_directory, numerical_attributes, categorical_attributes, seed,
        result, all_observed_arrivals, all_observed_departures, train_arrivals,
        test_arrivals, train_departures, test_departures, create_plots, record_candidates,
        cuemin_waiting_areas):
    cuemin = CueMin(
        seed_for_distributions=seed,
        categorical_attribute_names=categorical_attributes,
        numerical_attribute_names=numerical_attributes,
        nr_of_load_clusters_candidates=[], verbose=False,
        waiting_area_candidates=cuemin_waiting_areas)
    inferred_model, result_dict = evaluate_queue_inference(
        all_observed_arrivals, all_observed_departures,
        train_arrivals, test_arrivals, train_departures, test_departures,
        cuemin, 'CueMin',
        ground_truth_model, i,
        nr_timesteps_for_simulation, test_observation_time,
        experiment_directory, create_plots
    )
    result.append(
        result_dict | create_departure_time_predictor_dict(
            QueueDepartureTimePredictor(inferred_model, seed=seed),
            all_observed_arrivals, train_departures, test_departures
        )
    )
    if record_candidates:
        directory = os.path.join(experiment_directory, 'CueMin')
        if not os.path.exists(directory):
            os.mkdir(directory)
        cuemin.get_recording_dataframe().to_csv(
            os.path.join(experiment_directory, 'CueMin', f'cuemin_{i}.csv'), index=False, sep=';')

    linear_cuemin = CueMin(
        seed_for_distributions=seed,
        categorical_attribute_names=categorical_attributes,
        numerical_attribute_names=numerical_attributes,
        nr_of_load_clusters_candidates=[], verbose=False,
        waiting_area_candidates=cuemin_waiting_areas, search_strategy_name='linear',
        patience=10)
    inferred_model, result_dict = evaluate_queue_inference(
        all_observed_arrivals, all_observed_departures,
        train_arrivals, test_arrivals, train_departures, test_departures,
        linear_cuemin, 'LinearCueMin',
        ground_truth_model, i,
        nr_timesteps_for_simulation, test_observation_time,
        experiment_directory, create_plots
    )
    result.append(
        result_dict | create_departure_time_predictor_dict(
            QueueDepartureTimePredictor(inferred_model, seed=seed),
            all_observed_arrivals, train_departures, test_departures
        )
    )

    # sa_cuemin = CueMin(
    #     seed_for_distributions=seed,
    #     categorical_attribute_names=categorical_attributes,
    #     numerical_attribute_names=numerical_attributes,
    #     nr_of_load_clusters_candidates=[], verbose=False,
    #     waiting_area_candidates=cuemin_waiting_areas, search_strategy_name='sa-100')
    # inferred_model, result_dict = evaluate_queue_inference(
    #     all_observed_arrivals, all_observed_departures,
    #     train_arrivals, test_arrivals, train_departures, test_departures,
    #     sa_cuemin, 'SA-CueMin',
    #     ground_truth_model, i,
    #     nr_timesteps_for_simulation, test_observation_time,
    #     experiment_directory, create_plots
    # )
    # result.append(
    #     result_dict | create_departure_time_predictor_dict(
    #         QueueDepartureTimePredictor(inferred_model, seed=seed),
    #         all_observed_arrivals, train_departures, test_departures
    #     )
    # )

    if len(cuemin_waiting_areas) > 1:
        if 'PQ' in ground_truth_model.get_waiting_area().get_discipline_name():
            informed_waiting_area = 'PQ-c'
        elif 'LCFS' in ground_truth_model.get_waiting_area().get_discipline_name():
            informed_waiting_area = 'LCFS'
        elif 'FCFS' in ground_truth_model.get_waiting_area().get_discipline_name():
            informed_waiting_area = 'FCFS'
        else:
            return
        informed_cuemin = CueMin(
            seed_for_distributions=seed,
            categorical_attribute_names=categorical_attributes,
            numerical_attribute_names=numerical_attributes,
            nr_of_load_clusters_candidates=[], verbose=False,
            waiting_area_candidates=set([informed_waiting_area]))
        inferred_model, result_dict = evaluate_queue_inference(
            all_observed_arrivals, all_observed_departures,
            train_arrivals, test_arrivals, train_departures, test_departures,
            informed_cuemin, 'DInformedCueMin',
            ground_truth_model, i,
            nr_timesteps_for_simulation, test_observation_time,
            experiment_directory, create_plots
        )
        result.append(
            result_dict | create_departure_time_predictor_dict(
                QueueDepartureTimePredictor(inferred_model, seed=seed),
                all_observed_arrivals, train_departures, test_departures
            )
        )

def create_departure_time_predictor_dict(
        predictor: DepartureTimePredictor,
        all_observed_arrivals: list[tuple[Job, int]],
        train_departures: list[tuple[Job, int]],
        test_departures: list[tuple[Job, int]]) -> dict[str, str|int|float]:
    predicted_departure_times = predictor.predict(all_observed_arrivals)

    return {
        'MAE_departure_time_train': Statistics(abs(
                actual_departure_time - predicted_departure_times[job]
            ) for job, actual_departure_time in train_departures).mean(),
        'MAE_departure_time_test': Statistics(abs(
                actual_departure_time - predicted_departure_times[job]
            ) for job, actual_departure_time in test_departures).mean(),
        'MAPE_departure_time_train': Statistics(abs(
                (actual_departure_time - predicted_departure_times[job]) / actual_departure_time
            ) for job, actual_departure_time in train_departures).mean(),
        'MAPE_departure_time_test': Statistics(abs(
                (actual_departure_time - predicted_departure_times[job]) / actual_departure_time
            ) for job, actual_departure_time in test_departures).mean(),
    }

def create_data_from_executed_ground_truth_model(ground_truth_model: Queue, test_observation_time: float):
    observed_arrivals = [
        (job, arrival_time) for job, arrival_time in zip(
            ground_truth_model.get_arrival_process().get_recorded_jobs(),
            ground_truth_model.get_arrival_process().get_recorded_arrival_times()
        )
    ]
    observed_departures = [
        (job, exit_time) for job, exit_time in zip(
            *ground_truth_model.get_exit().get_recording()
        )
    ]
    train_arrivals = [
        (job, arrival_time) for job, arrival_time in observed_arrivals
        if arrival_time < test_observation_time
    ]
    train_departures = [
        (job, exit_time) for job, exit_time in observed_departures
        if exit_time < test_observation_time
    ]
    test_arrivals = [
        (job, arrival_time) for job, arrival_time in observed_arrivals
        if arrival_time >= test_observation_time
    ]
    test_departures = [
        (job, exit_time) for job, exit_time in observed_departures
        if exit_time >= test_observation_time
    ]
    return (
        observed_arrivals, observed_departures,
        train_arrivals, train_departures,
        test_arrivals, test_departures
    )

def run_model(
        model: Queue, nr_timesteps_for_simulation: int, test_observation_time: int):
    model.set_waiting_time_observer(WaitingTimeRecordingObserver(test_observation_time))
    model.set_sojourn_time_observer(SojournTimeRecordingObserver(test_observation_time))
    model.set_queue_length_observer(QueueLengthRecordingObserver(test_observation_time))
    model.set_exit(ListCollectorExit())

    environment = Environment(verbose=False)
    model.schedule_next_arrival(environment)

    environment.run_timesteps(nr_timesteps_for_simulation)

def run_ground_truth_model(
        model: Queue, nr_timesteps_for_simulation: int,
        nr_of_jobs_per_server: int,
        test_observation_time_proportion: float,
        nr_of_changes_distribution_train=DiscreteDegenerateDistribution(0),
        nr_of_changes_distribution_test=DiscreteDegenerateDistribution(0),
        probability_add_server=0.5,
        probability_remove_server=0.5,
        random: Random = Random()) -> int:
    model.set_queue_length_observer(QueueLengthRecordingObserver())
    model.set_exit(ListCollectorExit())

    environment = Environment(verbose=False)
    model.schedule_next_arrival(environment)

    nr_of_change_operations_during_training = nr_of_changes_distribution_train.get_next_sample()
    test_observation_time = int(nr_timesteps_for_simulation * test_observation_time_proportion)
    for i in range(nr_of_change_operations_during_training):
        time_of_change = (
            test_observation_time
            // (nr_of_change_operations_during_training + 1) * (i + 1)
        )
        environment.schedule_event(
            np.random.default_rng(random.randint(1, 1000000)).choice(
                [AddServerEvent(time_of_change, model), RemoveServerEvent(time_of_change, model)],
                1,
                p=[probability_add_server, probability_remove_server]
            )[0]
        )

    nr_of_change_operations_during_test = nr_of_changes_distribution_test.get_next_sample()
    for i in range(nr_of_change_operations_during_test):
        time_of_change = (
            test_observation_time + (nr_timesteps_for_simulation - test_observation_time)
            // nr_of_change_operations_during_test * i
        )
        environment.schedule_event(
            np.random.default_rng(random.randint(1, 1000000)).choice(
                [AddServerEvent(time_of_change, model), RemoveServerEvent(time_of_change, model)],
                1,
                p=[probability_add_server, probability_remove_server]
            )[0]
        )

    if nr_of_jobs_per_server is not None:
        required_nr_of_jobs = nr_of_jobs_per_server * model.get_nr_of_servers()
        environment.run_until(lambda: len(model.get_exit()) >= required_nr_of_jobs or not environment.has_open_events())
    else:
        environment.run_timesteps(nr_timesteps_for_simulation)

    return environment.get_current_time()

def store_data_to_data_log_directory(
        observed_arrivals: list[tuple[Job, int]],
        observed_departures: list[tuple[Job, int]],
        i: int, data_log_directory: str):
    with open(os.path.join(data_log_directory, f'arrivals_{i}.txt'), 'w', encoding='utf-8') as f:
        for job, arrival_time in observed_arrivals:
            f.write(f'{arrival_time};{job.job_id}\n')
    with open(os.path.join(data_log_directory, f'departures_{i}.txt'), 'w', encoding='utf-8') as f:
        for job, departure_time in observed_departures:
            f.write(f'{departure_time};{job.job_id}\n')
    with open(os.path.join(data_log_directory, f'jobs_{i}.txt'), 'w', encoding='utf-8') as f:
        for job in set(chain(
                (j for j,_ in observed_arrivals),
                (j for j,_ in observed_departures))):
            f.write(str(job))
            f.write('\n')

def create_model_evaluation_dict(
        model: Queue, model_name: str, i: int,
        nr_timesteps_for_simulation: float, test_observation_time: float,
        train_arrivals: list[tuple[Job, int]], test_arrivals: list[tuple[Job, int]],
        train_departures: list[tuple[Job, int]], test_departures: list[tuple[Job, int]],
        departure_time_per_job: dict[Job, int],
        inference_time_in_s: int = None) -> dict:
    return {
        'i': i,
        'total_timesteps': nr_timesteps_for_simulation,
        'train_timesteps': test_observation_time,
        'A': model.get_arrival_process(),
        'model': model_name,
        'inference_time_in_s': inference_time_in_s,
        'D': model.get_waiting_area().get_discipline_name(),
        'B': str(model.get_batch_size_distribution()),
        'S': model.get_service_time_name(),
        'c': model.get_nr_of_servers(),
        'nr_of_processed_jobs': len(model.get_exit()),
        'mean_service_batch_size': model.get_batch_size_distribution().get_mean(),
        'mean_interarrival_time': model.get_arrival_process().get_mean_interarrival_time(),
        'mean_waiting_time': model.get_waiting_time_observer().get_mean_waiting_time(),
        'median_waiting_time': model.get_waiting_time_observer().get_median_waiting_time(),
        'max_waiting_time': model.get_waiting_time_observer().get_max_waiting_time(),
        'mean_sojourn_time': model.get_sojourn_time_observer().get_mean_sojourn_time(),
        'median_sojourn_time': model.get_sojourn_time_observer().get_median_sojourn_time(),
        'max_sojourn_time': model.get_sojourn_time_observer().get_max_sojourn_time(),
        'max_queue_length': model.get_queue_length_observer().get_max_queue_length(),
        'MDL_train': compute_mdl(model, train_arrivals, train_departures, departure_time_per_job, 1),
        'MDL_all': compute_mdl(
            model, train_arrivals + test_arrivals,
            train_departures + test_departures, departure_time_per_job, 1)
    }

def plot_model_execution(model: Queue, model_name: str, ground_truth_model_index: int, experiment_directory: str):
    directory = os.path.join(experiment_directory, model_name)
    sojourn_time_directory = os.path.join(directory, 'soujourn time')
    if not os.path.exists(directory):
        try:
            os.mkdir(directory)
            os.mkdir(sojourn_time_directory)
        except FileExistsError:
            pass
    if model.get_sojourn_time_observer().get_sojourn_times():
        with PlotContext(
                'Sojourn Time Histogram', show=False,
                filepath=os.path.join(
                    sojourn_time_directory, f'sojourn_time_histogram_{ground_truth_model_index}'
                )) as plot_context:
            plot_histogram(
                model.get_sojourn_time_observer().get_sojourn_times(),
                ax=plot_context.get_axes())

def evaluate_queue_inference(
        observed_arrivals: list[tuple[Job, int]],
        observed_departures: list[tuple[Job, int]],
        train_arrivals: list[tuple[Job, int]],
        test_arrivals: list[tuple[Job, int]],
        train_departures: list[tuple[Job, int]],
        test_departures: list[tuple[Job, int]],
        queue_inference: QueueMiner, model_name: str,
        ground_truth_model: Queue, ground_truth_model_index: int,
        nr_timesteps_for_simulation: int, test_observation_time: int,
        experiment_directory: str, create_plots: bool) -> tuple[Queue|None, Union[dict, str]|str]:
    try:
        stopwatch = Stopwatch()
        stopwatch.start()
        inferred_model = queue_inference.infer_queue(train_arrivals, train_departures)
        inference_time_in_s = stopwatch.get_elapsed_time().total_seconds()
        copy_of_inferred_model = inferred_model.copy()
        copy_of_inferred_model.set_arrival_process(
            RecordingArrival(ground_truth_model.get_arrival_process().to_fixed_arrival_process()))
        run_model(copy_of_inferred_model, nr_timesteps_for_simulation, test_observation_time)
        if create_plots:
            plot_model_execution(copy_of_inferred_model, model_name, ground_truth_model_index, experiment_directory)
        return inferred_model, create_model_evaluation_dict(
            copy_of_inferred_model, model_name, ground_truth_model_index,
            nr_timesteps_for_simulation, test_observation_time,
            train_arrivals, test_arrivals,
            train_departures, test_departures,
            dict(observed_departures),
            inference_time_in_s=inference_time_in_s)
    # except Exception as e:
    except ConnectionError as e:
        return None, f'no result for {model_name} on {ground_truth_model_index}-th dataset because of {e}'

def evaluate_batch_miner(
        observed_arrivals: list[tuple[Job, float]],
        observed_departures: list[tuple[Job, float]],
        batch_miner: PerformanceSpectrumBatchMiner, miner_name: str,
        ground_truth_model_index: int,
        nr_timesteps_for_simulation: float, test_observation_time: float):
    batches = batch_miner.group_batches(observed_arrivals, observed_departures)
    batch_size_statistics = Statistics(len(b) for b in batches)
    return {
        'i': ground_truth_model_index,
        'total_timesteps': nr_timesteps_for_simulation,
        'train_timesteps': test_observation_time,
        'model': miner_name,
        'mean_service_batch_size': batch_size_statistics.mean(),
    }

def evaluate_waiting_area_estimator(
        observed_arrivals: list[tuple[Job, float]],
        observed_departures: list[tuple[Job, float]],
        waiting_area_miner: WaitingAreaEstimator, miner_name: str,
        ground_truth_model_index: int,
        nr_timesteps_for_simulation: float, test_observation_time: float):
    waiting_area = waiting_area_miner.infer_waiting_area(observed_arrivals, observed_departures)
    return {
        'i': ground_truth_model_index,
        'total_timesteps': nr_timesteps_for_simulation,
        'train_timesteps': test_observation_time,
        'model': miner_name,
        'D': waiting_area.get_discipline_name()
    }

def evaluate_nr_of_servers_estimator(
        observed_arrivals: list[tuple[Job, float]],
        observed_departures: list[tuple[Job, float]],
        estimator: NrOfServersEstimator, miner_name: str,
        ground_truth_model_index: int,
        nr_timesteps_for_simulation: float, test_observation_time: float):
    return {
        'i': ground_truth_model_index,
        'total_timesteps': nr_timesteps_for_simulation,
        'train_timesteps': test_observation_time,
        'model': miner_name,
        'c': estimator.estimate_nr_of_servers(observed_arrivals, observed_departures),
    }

class RunReconstructionSingleQueueExperiment:

    def __init__(
            self,
            population: InfinitePopulation,
            arrival_distributions: list[Distribution],
            arrival_batch_size_distributions: list[Distribution],
            waiting_area_factories: list[Callable[[],WaitingArea]],
            batch_size_distributions: list[Distribution],
            service_time_distributions: list[Distribution],
            nr_of_servers: list[int],
            experiment_name: str = 'experiment',
            nr_timesteps_for_simulation: int = 1000,
            nr_of_jobs_per_server: int|None = None,
            max_nr_of_arriving_jobs: int = -1,
            max_nr_of_models: int = -1,
            copies_per_stochastic_model: int = 0,
            computation_engine: ComputationEngine = SingleThreadComputationEngine(),
            create_plots: bool = False,
            log_data: bool = False,
            should_down_sample_models_per_nr_of_servers: bool = False,
            nr_of_changes_distribution_train=DiscreteDegenerateDistribution(0),
            nr_of_changes_distribution_test=DiscreteDegenerateDistribution(0),
            probability_add_server=0.5,
            probability_remove_server=0.5,
            record_cuemin_candidates: bool=False,
            cuemin_waiting_areas: list[str]=['LCFS', 'FCFS', 'PQ-c'],
            seed: int = None):
        if nr_of_jobs_per_server is not None:
            experiment_name = f'{experiment_name} - {nr_of_jobs_per_server} nr_of_jobs_per_server'
            if nr_of_changes_distribution_train != DiscreteDegenerateDistribution(0) \
            or nr_of_changes_distribution_test != DiscreteDegenerateDistribution(0):
                raise NotImplementedError('cannot set change distribution if nr_of_jobs_per_server is set')
        else:
            experiment_name = f'{experiment_name} - {nr_timesteps_for_simulation} timesteps'
        self.logger = ExperimentLogger(
            'experiments/evaluation/results/reconstruction_single_queue',
            experiment_name)
        self.__data_log_directory = os.path.join(self.logger.get_path_to_experiment_directory(), 'data')
        os.mkdir(self.__data_log_directory)
        self.__evaluation_csv_file = self.logger.open_text_file('evaluation.csv')
        self.__evaluation_csv_writer = csv.DictWriter(
            self.__evaluation_csv_file,
            fieldnames=[
                'i', 'total_timesteps', 'train_timesteps',
                'A', 'model', 'inference_time_in_s', 'D', 'B', 'S', 'mean_service_batch_size', 'c',
                'nr_of_processed_jobs', 'mean_interarrival_time',
                'mean_waiting_time', 'median_waiting_time', 'max_waiting_time',
                'mean_sojourn_time', 'median_sojourn_time', 'max_sojourn_time',
                'max_queue_length', 'MAE_departure_time_train', 'MAE_departure_time_test',
                'MAPE_departure_time_test', 'MAPE_departure_time_train',
                'MDL_train', 'MDL_all'
            ], delimiter=';')
        self.__evaluation_csv_writer.writeheader()
        self.__seed = seed
        self.__arrival_distributions = arrival_distributions
        self.__max_nr_of_arriving_jobs = max_nr_of_arriving_jobs
        self.__nr_timesteps_for_simulation = nr_timesteps_for_simulation
        self.__max_nr_of_models = max_nr_of_models
        self.__test_observation_time_proportion = 0.8
        self.__nr_of_jobs_per_server = nr_of_jobs_per_server
        self.__waiting_area_factories = waiting_area_factories
        self.__arrival_batch_size_distributions = arrival_batch_size_distributions
        self.__batch_size_distributions = batch_size_distributions
        self.__service_time_distributions = service_time_distributions
        self.__nr_of_servers = nr_of_servers
        self.__population = population
        self.__computation_engine = computation_engine
        self.__create_plots = create_plots
        self.__log_data = log_data
        self.__should_down_sample_models_per_nr_of_servers = should_down_sample_models_per_nr_of_servers
        self.__nr_of_changes_distribution_train = nr_of_changes_distribution_train
        self.__nr_of_changes_distribution_test = nr_of_changes_distribution_test
        self.__probability_add_server = probability_add_server
        self.__probability_remove_server = probability_remove_server
        self.__copies_per_stochastic_model = copies_per_stochastic_model
        self.__record_cuemin_candidates = record_cuemin_candidates
        self.__cuemin_waiting_areas = cuemin_waiting_areas

    def __create_ground_truth_models(self) -> list[Queue]:
        ground_truth_models = []

        possible_arrivals = []
        for arrival_distribution in self.__arrival_distributions:
            arrival = ArrivalWithDistribution(self.__population, arrival_distribution)
            if self.__max_nr_of_arriving_jobs > 0:
                arrival = LimitedArrival(arrival, self.__max_nr_of_arriving_jobs)
            for batch_size_distribution in self.__arrival_batch_size_distributions:
                if isinstance(batch_size_distribution, DiscreteDegenerateDistribution) \
                and batch_size_distribution.get_mean() == 1:
                    possible_arrivals.append(arrival)
                else:
                    possible_arrivals.append(BatchArrivalAdapter(
                        arrival.copy(), batch_size_distribution))
        possible_servers = [
            [
                Server(ServiceTimeWithDistribution(service_time_distribution.copy()))
                for _ in range(nr_of_servers)
            ]
            for service_time_distribution, nr_of_servers in product(
                self.__service_time_distributions, self.__nr_of_servers)
        ]

        for arrival_process, servers, waiting_area_factory, batch_size_distribution in product(
                possible_arrivals, possible_servers, self.__waiting_area_factories,
                self.__batch_size_distributions):
            ground_truth_models.append(
                Queue(
                    RecordingArrival(arrival_process),
                    servers,
                    waiting_area=waiting_area_factory(),
                    batch_size_distribution=batch_size_distribution,
                    exit_point=ListCollectorExit()
                )
            )

        ground_truth_models = [
            model for model in ground_truth_models
            if self.__model_is_estimated_to_be_identifiable(model)
        ]
        nr_of_models_per_nr_of_servers = defaultdict(int)
        models_per_nr_of_servers = defaultdict(list)
        for model in ground_truth_models:
            nr_of_models_per_nr_of_servers[model.get_nr_of_servers()] += 1
            models_per_nr_of_servers[model.get_nr_of_servers()].append(model)
        self.logger.log(f'distribution of models per number of servers: {nr_of_models_per_nr_of_servers}')
        if self.__should_down_sample_models_per_nr_of_servers:
            ground_truth_models = self.__down_sample_models_per_nr_of_servers(models_per_nr_of_servers)

        if self.__copies_per_stochastic_model > 0:
            extended_ground_truth_models = []
            for model in ground_truth_models:
                if not model.get_servers()[0].get_service_time_definition().is_deterministic() \
                or not model.get_batch_size_distribution().is_deterministic():
                    for i in range(self.__copies_per_stochastic_model):
                        model_copy = model.copy()
                        if self.__seed is not None:
                            model_copy.set_seed(self.__seed + i + 1)
                        extended_ground_truth_models.append(model_copy)
            ground_truth_models.extend(extended_ground_truth_models)

        if self.__max_nr_of_models > 0 and len(ground_truth_models) > self.__max_nr_of_models:
            Random(self.__seed).shuffle(ground_truth_models)
            ground_truth_models = ground_truth_models[:self.__max_nr_of_models]

        return ground_truth_models

    def __down_sample_models_per_nr_of_servers(self, models_per_nr_of_servers: dict[int, list[Queue]]) -> list[Queue]:
        random = Random(self.__seed)
        min_nr_of_models = len(min(models_per_nr_of_servers.values(), key=len))
        selected_models = []
        for model_list in models_per_nr_of_servers.values():
            selected_models.extend(random.sample(model_list, k=min_nr_of_models))
        return selected_models

    def __model_is_estimated_to_be_identifiable(self, model: Queue) -> bool:
        #use the erlang-c formula to estimate that the waiting time of a customer is non-zero
        #if there is waiting time, the servers are all used and we have a realistic chance
        #to detect the ground-truth model. of course, the erlang-c formula does only
        #strictly apply to M/M/c models, i.e. we approximate the ground-truth model
        #with a M/M/c model to decide its identifiability
        arrival_rate = model.get_arrival_process().get_mean_arrival_rate()
        service_rate = (
            model.get_batch_size_distribution().get_mean() /
            model.get_servers()[0].get_service_time_definition().get_distribution().get_mean()
        )
        a = arrival_rate / service_rate
        b = pow(a, model.get_nr_of_servers()) / factorial(model.get_nr_of_servers())
        try:
            b *= model.get_nr_of_servers() / (model.get_nr_of_servers() - a)
        except ZeroDivisionError:
            return True
        d = 0
        for n in range(0, model.get_nr_of_servers()):
            d += pow(a, n) / factorial(n)
        P_1 = b / (b + d)
        #return True if the probability of waiting time beging greater than 0 is >> 0
        return P_1 > 0.001

    def run(self):
        ground_truth_model_list = self.__create_ground_truth_models()
        ground_truth_model_list = list(enumerate(ground_truth_model_list))
        self.logger.log(f'generated {len(ground_truth_model_list)} ground-truth models')
        ground_truth_model_list = self.__computation_engine.create_partitionable_list(ground_truth_model_list)
        for csv_row_list in ground_truth_model_list.map({
                    'nr_timesteps_for_simulation': self.__nr_timesteps_for_simulation,
                    'nr_of_jobs_per_server': self.__nr_of_jobs_per_server,
                    'test_observation_time_proportion': self.__test_observation_time_proportion,
                    'data_log_directory': self.__data_log_directory,
                    'experiment_directory': self.logger.get_path_to_experiment_directory(),
                    'numerical_attributes': self.__population.get_numerical_feature_names(),
                    'categorical_attributes': self.__population.get_categorical_feature_names(),
                    'log_data': self.__log_data,
                    'create_plots': self.__create_plots,
                    'nr_of_changes_distribution_train': self.__nr_of_changes_distribution_train,
                    'nr_of_changes_distribution_test': self.__nr_of_changes_distribution_test,
                    'probability_add_server': self.__probability_add_server,
                    'probability_remove_server': self.__probability_remove_server,
                    'record_cuemin_candidates': self.__record_cuemin_candidates,
                    'cuemin_waiting_areas': self.__cuemin_waiting_areas,
                    'seed': self.__seed
                }, run_evaluation_on_ith_ground_truth_model):
            for csv_row in csv_row_list:
                if isinstance(csv_row, dict):
                    self.__evaluation_csv_writer.writerow(csv_row)
                else:
                    self.logger.log(csv_row)
        self.__evaluation_csv_file.close()
        pd.read_csv(
            self.logger.get_absolute_path('evaluation.csv'), delimiter=';'
        ).to_excel(self.logger.get_absolute_path('evaluation.xlsx'))
        self.logger.log('finished experiment')
        self.logger.close()

def parse_waiting_area_factories() -> list[Callable[[],WaitingArea]]:
    categories = list(range(arguments.nr_of_categories))
    random_generator = Random(arguments.seed)
    def generate_logistic_regression() -> LogisticRegression:
        lr = LogisticRegression()
        # 1 + ... because of "difference of arrival time" as additional feature
        lr.coef_ = np.array([[
            0 if random_generator.choice([True, False]) else random_generator.random()
            for _ in range(
                1 + 2 * (
                    arguments.nr_of_numerical_features +
                    arguments.nr_of_categorical_features * arguments.nr_of_categories
                )
            )
        ]])
        lr.classes_ = np.array([True, False])
        lr.intercept_ = 0
        return lr
    waiting_areas = {
        'FCFS': FirstComeFirstServeWaitingArea,
        'LCFS': LastComeFirstServeWaitingArea,
        'PQc-FCFS': lambda: PriorityClassWaitingArea(
            random_generator.choice(population.get_categorical_feature_names()),
            categories, FirstComeFirstServeWaitingArea),
        'PQc-LCFS': lambda: PriorityClassWaitingArea(
            random_generator.choice(population.get_categorical_feature_names()),
            categories, LastComeFirstServeWaitingArea),
        'PQ-LR': lambda: PairwisePriorityClassifierWaitingArea(SklearnPairwisePriorityClassifier(
            generate_logistic_regression(),
            JobToVectorTransformer(
                population.get_numerical_feature_names(),
                population.get_categorical_feature_names(),
                {}, {f: OneHotEncoder(categories) for f in population.get_categorical_feature_names()}))
        )
    }
    return [waiting_areas[s] for s in arguments.D]

def parse_distributions(distribution_list: list[str]) -> list[Distribution]:
    parsed_distribution_list = []
    for distribution in distribution_list:
        parsed_distribution_list.append(parse_distribution(distribution, seed=arguments.seed))
    return parsed_distribution_list

def parse_nr_of_servers(c_list: list[str]) -> list[int]:
    parsed_c_list = []
    for c_str in c_list:
        if ':' in c_str:
            min_c_str, max_c_str = c_str.split(':')
            parsed_c_list.extend(range(int(min_c_str), int(max_c_str) + 1))
        else:
            parsed_c_list.append(int(c_str))
    return parsed_c_list

def parse_computation_engine() -> ComputationEngine:
    if arguments.execution == 'S':
        return SingleThreadComputationEngine()
    elif arguments.execution.startswith('P'):
        return MultiprocessComputationEngine(int(arguments.execution[1:]))
    else:
        raise NotImplementedError(f'unknown execution config "{arguments.execution}"')

if __name__ == '__main__':
    parser = argparse.ArgumentParser('queue mining experiments on artificially generated single queues')

    parser.add_argument(
        '--experiment_name', type=str,
        help='name of the experiment',
        required=False, default='experiment')

    parser.add_argument(
        '--seed', type=int,
        help='random generator seed',
        required=False, default=None)

    parser.add_argument(
        '--nr_of_categorical_features', type=int,
        required=False, default=3)
    parser.add_argument(
        '--nr_of_categories', type=int,
        required=False, default=3)
    parser.add_argument(
        '--nr_of_numerical_features', type=int,
        required=False, default=3)

    parser.add_argument(
        '--nr_timesteps_for_simulation', type=int,
        help='total number of timesteps to simulate',
        required=False, default=10000)
    parser.add_argument(
        '--nr_of_jobs_per_server', type=int,
        help=(
            'Overrides nr_timesteps_for_simulation, such that the simulation runs, '
            'until this number of jobs per server have exited the system. '
            'This does not ensure that each server is used!'
        ),
        required=False, default=None)
    parser.add_argument(
        '--max_nr_of_arriving_jobs', type=int,
        help='limits the number of arrived jobs during the simulation. by default there is no limit',
        required=False, default=-1)

    parser.add_argument(
        '--max_nr_of_models', type=int,
        help='maximum number of ground-truth models to evaluate. default means no maximum.',
        required=False, default=-1)
    parser.add_argument(
        '--copies_per_stochastic_model', type=int,
        help='number of times a stochastic ground-truth model is duplicated.',
        required=False, default=0)

    parser.add_argument(
        '--D', type=str,
        help='possible waiting areas', nargs='+',
        required=False, default=['FCFS', 'LCFS', 'PQc-FCFS', 'PQc-LCFS'],
        choices=['FCFS', 'LCFS', 'PQc-FCFS', 'PQc-LCFS', 'PQ-LR'])

    parser.add_argument(
        '--c', type=str,
        help='determines the possible number of servers for the generated ground-truth models',
        nargs='+', required=False, default=['1:5'])

    parser.add_argument(
        '--A', type=str,
        help=('determines which inter-arrival distributions are possible. '
              'Dx denotes a degenerate (fixed) distribution of interarrival time size x, e.g. "D1". '
              'G-x-y denotes a Gamma distribution with shape x and rate y, e.g. "G-2-0.5".'),
        nargs='+', required=False, default=['D1'])

    parser.add_argument(
        '--S', type=str,
        help=('determines which service time distributions are possible. '
              'Dx denotes a degenerate (fixed) distribution of service time x, e.g. "D1". '
              'Ex denotes an exponential distribution with rate x, e.g. "E0.1". '
              'G-x-y denotes a Gamma distribution with shape x and rate y, e.g. "NB-2-0.5".'),
        nargs='+', required=False, default=['Geo0.1'])

    parser.add_argument(
        '--BS', type=str,
        help=('determines which batch size distributions for service are possible. '
              'Dx denotes a degenerate (fixed) distribution of batch size x, e.g. "D1". '
              'Px denotes a Poisson distribution with mean x, e.g. "P3".'),
        nargs='+', required=False, default=['D1'])
    parser.add_argument(
        '--BA', type=str,
        help=('determines which batch size distributions for arrival are possible. '
              'Dx denotes a degenerate (fixed) distribution of batch size x, e.g. "D1". '
              'Px denotes a Poisson distribution with mean x, e.g. "P3".'),
        nargs='+', required=False, default=['D1'])

    parser.add_argument(
        '--execution', type=str,
        help=('enables parallel execution to speed up the experiments '
              'S denotes single-threaded execution and is the default. '
              'Px denotes multiprocessing using x cores, e.g. "P3".'),
        required=False, default='S')

    parser.add_argument(
        '--create_plots', action='store_true',
        help=('add this flag to enable plotting of some data and model characteristics (e.g. sojourn times)'),
        required=False, default=False)

    parser.add_argument(
        '--record_cuemin_candidates', action='store_true',
        help=('add this flag to enable recording of cuemin search candidates. helpful for debugging'),
        required=False, default=False)
    parser.add_argument(
        '--cuemin_waiting_areas', type=str,
        help='possible waiting areas', nargs='+',
        required=False, default=['FCFS', 'LCFS', 'PQ-c'],
        choices=['FCFS', 'LCFS', 'PQ-c'])

    parser.add_argument(
        '--log_data', action='store_true',
        help=('set this to store the generated data into text files - useful for reproducibility'),
        required=False, default=False)

    parser.add_argument(
        '--down_sample_models_per_nr_of_servers', action='store_true',
        help=('set this to balance the number of models per nr of servers'),
        required=False, default=False)

    parser.add_argument(
        '--nr_of_changes_distribution_train', type=str,
        help='distribution for the number of change operations over time on ground-truth model during train time',
        required=False, default='D0')
    parser.add_argument(
        '--nr_of_changes_distribution_test', type=str,
        help='distribution for the number of change operations over time on ground-truth model during test time',
        required=False, default='D0')
    parser.add_argument(
        '--probability_add_server', type=float,
        required=False, default=0.5)
    parser.add_argument(
        '--probability_remove_server', type=float,
        required=False, default=0.5)

    arguments = parser.parse_args()

    population = parse_population(
        arguments.nr_of_categorical_features,
        arguments.nr_of_categories,
        arguments.nr_of_numerical_features,
        seed=arguments.seed
    )

    experiment = RunReconstructionSingleQueueExperiment(
        population,
        parse_distributions(arguments.A),
        parse_distributions(arguments.BA),
        parse_waiting_area_factories(),
        parse_distributions(arguments.BS),
        parse_distributions(arguments.S),
        parse_nr_of_servers(arguments.c),
        experiment_name=arguments.experiment_name,
        seed=arguments.seed,
        nr_timesteps_for_simulation=arguments.nr_timesteps_for_simulation,
        max_nr_of_arriving_jobs=arguments.max_nr_of_arriving_jobs,
        nr_of_jobs_per_server=arguments.nr_of_jobs_per_server,
        max_nr_of_models=arguments.max_nr_of_models,
        create_plots=arguments.create_plots,
        log_data=arguments.log_data,
        should_down_sample_models_per_nr_of_servers=arguments.down_sample_models_per_nr_of_servers,
        nr_of_changes_distribution_train=parse_distribution(
            arguments.nr_of_changes_distribution_train, seed=arguments.seed),
        nr_of_changes_distribution_test=parse_distribution(
            arguments.nr_of_changes_distribution_test, seed=arguments.seed),
        probability_add_server=arguments.probability_add_server,
        probability_remove_server=arguments.probability_remove_server,
        copies_per_stochastic_model=arguments.copies_per_stochastic_model,
        computation_engine=parse_computation_engine(),
        record_cuemin_candidates=arguments.record_cuemin_candidates,
        cuemin_waiting_areas=arguments.cuemin_waiting_areas
    )
    with experiment.logger.open_text_file('arguments.json') as f:
        f.write(json.dumps(vars(arguments), indent=4, sort_keys=True))
    experiment.run()