import numpy as np
import multiprocessing as mp
from pathos.multiprocessing import ProcessingPool

from matplotlib import pyplot as plt
from nonlinear_hawkes.utils.postprocessing.multivariate_postprocessing import select_arrivals
from rnpp.utils.queues_utils.service_functions import infinitesimal_propagator, service_time_from_arrivals_and_departures
from rnpp.utils.queues_utils.postprocessing import create_deep_service_format


from spatio_temporal_hawkes.utils.support import uniform_points_in_box, gp_mesh

from rnpp.queues.datatypes import Services, Queues
from nonlinear_hawkes.gaussian_nonlinear_hawkes.gnh import nonlinear_gaussian_hawkes
from nonlinear_hawkes.utils.plots.arrivals_plot import multivariate_hawkes_arrivals_plot

from spatio_temporal_hawkes.utils.arrivals_plot import temporal_plot_simulation
from spatio_temporal_hawkes.utils.support import uniform_points_in_box, gp_mesh
from spatio_temporal_hawkes.gaussian_hawkes.gaussian_self_exciting import gaussian_self_exciting, sigmoidal


def prepare_parameters(arrivals_type="gaussian",service_type="processor_sharing"):
    temporal_support = [(0., 1000.)]
    if arrivals_type == "gaussian":
        """SUPPORT"""
        arrival_type = "hawkes_process"
        number_of_xy = (10,)
        temporal_inducing_points = gp_mesh(temporal_support, number_of_xy)["points"]

        """GP PARAMETERS"""
        lenght_scale_x = 100.
        kernel_parameters = {"gamma": 1.,
                             "sigma": 1.,
                             "lenght_scale": np.asarray([[1. / lenght_scale_x ** 2]]).tolist(),
                             "dimension": 1}

        spatial_excitations_parameters = {"type": "empty"}
        temporal_excitations_parameters = {"type": "exponential", "alpha": .2, "support": temporal_support}
        marks_parameters = {"type": "gaussian_process_magnitudes",
                            "type_of_marks": "temporal",
                            "kernel_parameters": kernel_parameters,
                            "number_of_xy": number_of_xy,
                            "inducing_points": temporal_inducing_points.tolist(),
                            "inducing_function": None,
                            "support": temporal_support}

        model_parameters = {"read_model": False,
                            "model_label": "gauss_self_exciting_1D",
                            "model_dir": "/home/anon/Desktop/Doctorado/Projects/GaussianHawkesProcess/Experiments/Gaussian_Self_Exciting/",
                            "spatial_support": None,  # [(0.,100.),(0.,100.)]
                            "temporal_support": temporal_support,
                            "lambda_0": 4.,
                            "base_intensity_type": "uniform",
                            "spatial_excitations_parameters": spatial_excitations_parameters,
                            "temporal_excitations_parameters": temporal_excitations_parameters,
                            "marks_parameters": marks_parameters}

    elif arrivals_type == "nonlinear_hawkes":
        arrival_type = "nonlinear_hawkes"

        excitations_parameters = {"type": "exponential",
                                  "alpha": 1.}

        model_parameters = {"temporal_support": temporal_support,
                            "number_of_processes": 2,
                            "lambda_0": 2.,
                            "mu_k_0": [10., 10.],
                            "connectivity_matrix": [[1., 0.],
                                                    [0., 1.]],
                            "excitations_parameters": excitations_parameters}

    if service_type == "processor_sharing":
        processor_type = "processor_sharing"
        service_parameters = {"service_type": processor_type, "service_time": 0.5}
        service_simulate_parameters = {"temporal_support": temporal_support}
    elif service_type == "phase_type_service":
        number_of_servers = 30
        processor_type = "phase_type_service"
        Q = infinitesimal_propagator(number_of_servers)
        service_parameters = {"infinitesimal_propagator": Q, "number_of_servers": number_of_servers}
        service_simulate_parameters = {"T": temporal_support[0][1]}

    queues_parameters = {"arrivals_parameters": {"arrivals_type": arrival_type,
                                                 "parameters": model_parameters},
                         "service_parameters": {"service_type": processor_type,
                                                "service_parameters": service_parameters}}

    queues_simulation_parameters = {"arrivals": model_parameters["temporal_support"],
                                    "service": service_simulate_parameters}

    return temporal_support, queues_parameters, queues_simulation_parameters


def get_attribute(QUEUES,parameters):
    return getattr(QUEUES,'simulate')(parameters)



def simulation_worker(parameters):
    queues_object = parameters[0]
    queues_simulation_parameters = parameters[1]
    return queues_object.simulate(queues_simulation_parameters)


if __name__=="__main__":
    parallel = True

    number_of_process = 5
    NUMBER_OF_BATCHES = 10
    # arrivals_type = "gaussian"
    arrivals_type = "nonlinear_hawkes"
    #service_type = "processor_sharing"
    results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/new_nonlinear/"
    #results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/Nonlinear_PS/"
    #results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/nonlinearhawkes_processor_sharing/"

    service_type = "phase_type_service"
    # results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentNeuralPointProcess/Queues/PreprocessData/PhaseType/"
    # results_dir = "/home/anon/Desktop/Doctorado/Projects/RecurrentPointProcesses/Queues/service_data/nonlinearhawkes_phase_type/"

    temporal_support, queues_parameters, queues_simulation_parameters = prepare_parameters(arrivals_type, service_type)

    if not parallel:
        Queue = Queues(queues_parameters)
        ALL_ARRIVALS_DEPARTURES = []
        ALL_SERVICE_TIMES = []
        for k in range(NUMBER_OF_BATCHES):
            ALL_ARRIVALS_DEPARTURES.append(Queue.simulate((queues_simulation_parameters,k)))
            ALL_SERVICE_TIMES.append(service_time_from_arrivals_and_departures(ALL_ARRIVALS_DEPARTURES[k]))
        create_deep_service_format(ALL_ARRIVALS_DEPARTURES, temporal_support, results_dir)
    else:
        pool = mp.Pool(number_of_process)
        # GENERATE ALL THE DIFFERENT OBJECTS
        OBJECTS_LIST = ProcessingPool(processes=number_of_process).map(Queues,[(queues_parameters,i) for i in range(NUMBER_OF_BATCHES)])
        PARAMETERS_LIST = [(queues_simulation_parameters,i) for i in range(NUMBER_OF_BATCHES)]
        # PARALLELIZATION
        ALL_ARRIVALS_DEPARTURES = ProcessingPool(processes=number_of_process).map(get_attribute, OBJECTS_LIST, PARAMETERS_LIST)

        create_deep_service_format(ALL_ARRIVALS_DEPARTURES, temporal_support, results_dir)

"""
def PARALLEL_INFERENCE(GPH_PARAMETERS, INFERENCE_PARAMETERS, number_of_workers=2):
    # GENERATE ALL THE DIFFERENT OBJECTS
    OBJECTS_LIST = ProcessingPool(processes=number_of_workers).map(GaussianPolyaHawkesProcess,
                                                                   [(gpp_parameters, i) for i, gpp_parameters in
                                                                    enumerate(GPH_PARAMETERS)])
    PARAMETERS_LIST = [(inference_parameters,) for i, inference_parameters in enumerate(INFERENCE_PARAMETERS)]

    # PARALLELIZATION
    RESULTS = ProcessingPool(processes=number_of_workers).map(get_attribute, OBJECTS_LIST, PARAMETERS_LIST)

    return OBJECTS_LIST, RESULTS
"""
