"""
implements ideas from

Discovering Queues from Event Logs with Varying Levels of Information
Arik Senderovich, Sander J.J. Leemans, Shahar Harel, Avigdor Gal,
Avishai Mandelbaum, and Wil M.P. van der Aalst
"""

from collections import defaultdict
import numpy as np
from shs_sklearn_barebones.cluster import BarebonesKMeans

from prolothar_queue_mining.model.job import Job
from prolothar_queue_mining.model.environment import Environment
from prolothar_queue_mining.model.event import Event
from prolothar_queue_mining.model.distribution import PhaseTypeDistribution
from prolothar_queue_mining.model.service_time import ServiceTimeWithDistribution

from prolothar_queue_mining.prediction.departure_time import LoadClusterDepartureTimePredictor

class Kphf():
    """
    infers predictors, e.g. for the departure time of jobs based on

    Discovering Queues from Event Logs with Varying Levels of Information
    Arik Senderovich, Sander J.J. Leemans, Shahar Harel, Avigdor Gal,
    Avishai Mandelbaum, and Wil M.P. van der Aalst
    """

    def __init__(self, k: int = 4, distribution_seed: int = None):
        self.__k = k
        self.__distribution_seed = distribution_seed

    def infer_departure_time_predictor(self, observed_arrivals: list[tuple[Job, int]],
                                       observed_departures: list[tuple[Job, int]]):
        nr_of_other_jobs_in_system = self.__compute_nr_of_other_jobs_in_system(
            observed_arrivals, observed_departures)

        arrival_time_per_job = {job: time for job, time in observed_arrivals}
        departure_time_per_job = {job: time for job, time in observed_departures}
        sojourn_time_per_job = {
            job: departure_time - arrival_time_per_job[job]
            for job, departure_time in departure_time_per_job.items()
            if job in arrival_time_per_job
        }
        kmeans = BarebonesKMeans(n_clusters=self.__k)
        sojourn_times_per_cluster = defaultdict(list)
        job_list, system_size_list = zip(*[(j,n) for j,n in nr_of_other_jobs_in_system.items()])
        for job, cluster_label in zip(job_list, kmeans.fit_predict(np.array(system_size_list).reshape(-1,1))):
            if job in sojourn_time_per_job:
                sojourn_times_per_cluster[cluster_label].append(sojourn_time_per_job[job])

        return LoadClusterDepartureTimePredictor(kmeans, {
            cluster_label: ServiceTimeWithDistribution(PhaseTypeDistribution.fit(
                sojourn_times, seed=self.__distribution_seed))
            for cluster_label, sojourn_times in sojourn_times_per_cluster.items()
        })

    def __compute_nr_of_other_jobs_in_system(
            self, observed_arrivals: list[tuple[Job, int]],
            observed_departures: list[tuple[Job, int]]) -> dict[Job, int]:
        nr_of_other_jobs_in_system = {}
        jobs_in_system = set()

        environment = Environment()
        for job, arrival_time in observed_arrivals:
            environment.schedule_event(ArrivalEvent(
                job, arrival_time, nr_of_other_jobs_in_system, jobs_in_system))
        for job, departure_time in observed_departures:
            environment.schedule_event(DepartureEvent(
                job, departure_time, jobs_in_system))
        environment.run_timesteps(observed_departures[-1][1])

        return nr_of_other_jobs_in_system

class ArrivalEvent(Event):

    def __init__(self, job: Job, arrival_time: int, nr_of_other_jobs_in_system, jobs_in_system):
        super().__init__(arrival_time, prio=1)
        self.__job = job
        self.__nr_of_other_jobs_in_system = nr_of_other_jobs_in_system
        self.__jobs_in_system = jobs_in_system

    def execute(self, environment: Environment):
        self.__nr_of_other_jobs_in_system[self.__job] = len(self.__jobs_in_system)
        self.__jobs_in_system.add(self.__job)

class DepartureEvent(Event):

    def __init__(self, job: Job, departure_time: int, jobs_in_system):
        super().__init__(departure_time)
        self.__job = job
        self.__jobs_in_system = jobs_in_system

    def execute(self, environment: Environment):
        self.__jobs_in_system.discard(self.__job)