'''
Created on Jan 9, 2018

@author: anon
'''

import operator
import random

from nonlinear_hawkes.utils.postprocessing.multivariate_postprocessing import select_arrivals
from nonlinear_hawkes.gaussian_nonlinear_hawkes.gnh import nonlinear_gaussian_hawkes
from spatio_temporal_hawkes.gaussian_hawkes.gaussian_self_exciting import gaussian_self_exciting


from rnpp.utils.queues_utils.service_functions import arrivals_and_departures
from rnpp.utils.queues_utils.service_functions import sample_from_rate, service_time_function

#from gaussian_hawkes.queues.utils import arrivals_and_departures
#from gaussian_hawkes.queues.utils import sample_from_rate, service_time_function

import numpy as np

class Queues:
    
    def __init__(self,parameters):
        #======================================
        # ARRIVALS
        #======================================
        queues_parameters = parameters[0]
        self.id = parameters[1]
        np.random.seed(self.id)

        arrivals_parameters = queues_parameters["arrivals_parameters"]
        #if arrivals_parameters["arrival_type"] == "poisson":
        #    self.ArrivalsProcess = PoissonProcess(arrivals_parameters["parameters"])
        #elif arrivals_parameters["arrival_type"] == "gaussian_poisson":
        #    self.ArrivalsProcess = gaussian_hawkes_process(arrivals_parameters["parameters"])
        if arrivals_parameters["arrivals_type"] == "hawkes_process":
            self.ArrivalsProcess = gaussian_self_exciting(arrivals_parameters["parameters"])
        elif  arrivals_parameters["arrivals_type"] == "nonlinear_hawkes":
            self.ArrivalsProcess = nonlinear_gaussian_hawkes(arrivals_parameters["parameters"])
        
        self.arrivals_parameters = arrivals_parameters
        
        #======================================
        # SERVICE
        #======================================
        self.Service = Services((queues_parameters["service_parameters"],self.id))
                
    def simulate(self,simulation_parameters):
        np.random.seed(simulation_parameters[1])
        arrivals_simulation_parameters = simulation_parameters[0]["arrivals"]
        service_simulation_parameters = simulation_parameters[0]["service"]
        ARRIVALS = self.ArrivalsProcess.simulate(arrivals_simulation_parameters)
        if self.arrivals_parameters["arrivals_type"] == "nonlinear_hawkes":
            ARRIVALS = np.asarray(select_arrivals(ARRIVALS)[0])
        elif self.arrivals_parameters["arrivals_type"] == "hawkes_process":
            ARRIVALS = ARRIVALS[:,0]
        number_of_arrivals = len(ARRIVALS)
        print number_of_arrivals
        DEPARTURES = self.Service.simulate(ARRIVALS, service_simulation_parameters)
        
        ARRIVALS_AND_DEPARTURES = arrivals_and_departures(ARRIVALS, DEPARTURES) 
         
        return ARRIVALS_AND_DEPARTURES

#===================================================================================
# SERVICE 
#===================================================================================

class Services:
    
    def __init__(self,parameters):
        self.type_of_service = parameters[0]["service_type"]
        self.parameters = parameters[0]["service_parameters"]
        self.id = parameters[1]
        np.random.seed(self.id)

    
    def simulate(self,ARRIVALS,simulation_parameters,ServiceState=None):
        print "Hey"
        if self.type_of_service == "phase_type_service":
            return self.PhaseTypeService(ARRIVALS)
        if self.type_of_service == "processor_sharing":
            return self.ProcessorSharing(ARRIVALS,simulation_parameters["temporal_support"],self.parameters["service_time"])
    
    def PhaseTimeDistribution(self,initial_state=0):
        """
        We modify:
            https://vknight.org/unpeudemath/code/2015/08/01/simulating_continuous_markov_chains.html
        """
        try:
            Q = self.parameters["infinitesimal_propagator"]
        except:
            print "Wrong PhaseTimeService parameters"
            raise Exception 
    
        Q = list(Q)  # In case a matrix is input
        number_of_states = len(Q)
        state_space = range(len(Q))  # Index the state space
        time_spent = {s:0 for s in state_space}  # Set up a dictionary to keep track of time
        clock = 0  # Keep track of the clock
        current_state = initial_state  # First state
        
        STATES = [current_state]
        while (current_state != number_of_states - 1):
            # Sample the transitions
            sojourn_times = [sample_from_rate(rate) for rate in Q[current_state][:current_state]]
            sojourn_times += [np.infty]  # An infinite sojourn to the same state
            sojourn_times += [sample_from_rate(rate) for rate in Q[current_state][current_state + 1:]]
            # Identify the next state
            next_state = min(state_space, key=lambda x: sojourn_times[x])
            STATES.append(next_state)
            
            sojourn = sojourn_times[next_state]
            clock += sojourn
            current_state = next_state  # Transition
            
        return clock
        
    def PhaseTypeService(self,ARRIVALS):
        """
        Here we follow  Queues Driven by Hawkes Process Andrew daw, Jamol Pender (2017)

        We modify:
        https://vknight.org/unpeudemath/code/2015/08/01/simulating_continuous_markov_chains.html
        
        Return
        ------
        Time it takes to leave the system
        """
        numberOfServers = self.parameters["number_of_servers"]

        
        numberOfArrivals = len(ARRIVALS)
        numberOfBatches = int(numberOfArrivals/numberOfServers)
        DEPARTURES = {i:0. for i in range(numberOfArrivals)}
        SERVICE_TIMES = [self.PhaseTimeDistribution() for a in ARRIVALS]
        SERVER_TIME = {k:0. for k in range(numberOfServers)}

        arrivals_indexes = range(numberOfArrivals)

        print "Starting Phase Type Service Simulation: "
        #INITIALIZATION<------------------------------------------------------------------
        
        server_batch_index = 0
        #server_batch_arrivals = [arrival for arrival in ARRIVALS[server_batch_index*numberOfServers:server_batch_index*numberOfServers+numberOfServers]]
        server_batch_times = [service for service in SERVICE_TIMES[server_batch_index * numberOfServers:server_batch_index * numberOfServers + numberOfServers]]

        #server index per arrival
        server_allocations = {i:i for i in arrivals_indexes[0:numberOfServers]}
        #current server time
        server_time = {k:0. for k in range(numberOfServers)}
        
        server_batch_index = 0
        
        #SIMULATION<-----------------------------------------------------------------------
        while server_batch_index <= numberOfBatches:
        
            #DISPATCHING (ASSIGNING DEPARTURE TIMES AND NEW SERVERS)
            for arrival_index, server_index in server_allocations.iteritems():
                arrival_server = server_allocations[arrival_index]
                departure_time = max(server_time[arrival_server],ARRIVALS[arrival_index]) + server_batch_times[server_index]
                DEPARTURES[arrival_index] += departure_time 
                server_time[server_index] = departure_time
                
            #====================
            # REALLOCATE
            #====================
            server_batch_index += 1
            server_batch_arrivals = [arrival for arrival in ARRIVALS[server_batch_index*numberOfServers:server_batch_index*numberOfServers+numberOfServers]]
            servers_ready = sorted(server_time.iteritems(),key=operator.itemgetter(1))
            server_allocations = {}
            for k,arrival_index in zip(range(numberOfServers),arrivals_indexes[server_batch_index*numberOfServers:server_batch_index*numberOfServers + numberOfServers]):
                server_allocations[arrival_index] = servers_ready[k][0]
            
        return DEPARTURES
    
    def ProcessorSharing(self,ARRIVALS,temporal_support=None,service_time=1.,alpha=10.):
        """
        Here we follow:
        http://www.columbia.edu/~ks20/4404-Sigman/4404-Notes-PS.pdf
        From Professor Karl Sigman's Lecture Notes on Simulation (Center for Applied Probability Columbia)
        """
        print "Starting Processor Sharing Simulation "

        ARRIVALS.sort()
        if temporal_support==None:
            print "Completion time set as last arrivals"
            temporal_support = [(0,ARRIVALS[-1])]
        
        number_of_arrivals = len(ARRIVALS)
        DEPARTURES = {}

        arrivals_index = 0
        arrival_time = ARRIVALS[arrivals_index]
        service_completion_time = np.infty

        remaining_service_times = []
        number_of_costumers_in_system = 0

        MAX_ITERATIONS = 10*number_of_arrivals
        NUMBER_OF_ITERATIONS = 0

        completed_arrivals = 0
        t = temporal_support[0][0]
        while t < temporal_support[0][1]:
            #=========================
            # ARRIVAL IS NEXT EVENT
            #=========================
            if arrival_time < service_completion_time:
                if number_of_costumers_in_system == 0: #OK!!!!!!!!!
                    number_of_costumers_in_system=1
                    service_time = service_time_function(service_time)
                    remaining_service_times.append((service_time,arrivals_index))        #HERE WE INCLUDE THE SERVICE TIME FOR THIS PARTICULAR ARRIVAL
                    t = arrival_time
                    arrivals_index += 1
                    try:
                        arrival_time = ARRIVALS[arrivals_index]
                    except:
                        arrival_time = np.infty 
                    service_completion_time = arrival_time + service_time  #<------------!!!
                elif number_of_costumers_in_system > 0:
                    w = alpha + ((arrival_time - t)/number_of_costumers_in_system)                       #CORRECTS WAITING TIME ACCORDING TO HOW MANY WERE THERE!!
                    remaining_service_times = map(lambda x: (max(x[0]-w,0),x[1]),remaining_service_times)
                    number_of_costumers_in_system = number_of_costumers_in_system+1
                    t = arrival_time
                    
                    service_time = service_time_function(service_time)
                    remaining_service_times.append((service_time,arrivals_index))         #HERE WE INCLUDE THE SERVICE TIME FOR THIS PARTICULAR ARRIVAL
                    remaining_service_times.sort()
                    
                    arrivals_index += 1
                    try:
                        arrival_time = ARRIVALS[arrivals_index]
                    except:
                        arrival_time = np.infty
                        
                    service_completion_time = t + (remaining_service_times[0][0]*number_of_costumers_in_system)                
            #============================
            # SERVICE COMPLETION IS NEXT
            #============================
            elif(service_completion_time < arrival_time):
                currently_out = remaining_service_times.pop(0)
                DEPARTURES[currently_out[1]]  = service_completion_time
                if ARRIVALS[currently_out[1]] > service_completion_time:
                    print "Arrival after departure"
                completed_arrivals+=1
                if completed_arrivals == number_of_arrivals:
                    break
                else:
                    number_of_costumers_in_system = number_of_costumers_in_system-1   
                    if number_of_costumers_in_system == 0: 
                        t = service_completion_time
                        service_completion_time = np.infty
                    else:
                        w = alpha + ((service_completion_time - t)/(number_of_costumers_in_system+1))  #CORRECTS WAITING TIME ACCORDING TO HOW MANY WERE THERE!!
                        remaining_service_times = map(lambda x: (max(x[0]-w,0),x[1]),remaining_service_times)
                        t=service_completion_time
                        service_completion_time = t + remaining_service_times[0][0]*number_of_costumers_in_system
            elif((service_completion_time == np.infty) and (arrival_time == np.infty)):
                break
                
            NUMBER_OF_ITERATIONS+=1
            if NUMBER_OF_ITERATIONS > MAX_ITERATIONS:
                break

            if NUMBER_OF_ITERATIONS % 500 == 0:
                print "Processor sharring iteration: ",t

            
        return DEPARTURES
        