#  Python 3
#  Author: Jonathan Goodman, goodman@cims.nyu.edu
#  Fall 2024
#  for the class Scientific Computing
#  https://math.nyu.edu/~goodman/teaching/ScientificComputing2024/ScientificComputing.html

#  Illustrate an event-driven simulation that uses a priority queue

# This file: EventDrivenSimulation.py,

import heapq
import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(seed=19)

lam = 20.     # arrival rate parameter
mu  = .5     # decay rate parameter

alpha = 1./lam    # scale parameter for arrivals
beta  = 1./mu     # scale parameter for decay times

Tmax = 50.        # simulate to this time

n_obs = 500       # plot n(t) at this many uniformly spaced times
t_obs = np.linspace(0., Tmax, n_obs)  # uniformly spaced obs times
n_vals = np.zeros(n_obs)              # will be values of n(t) at obs times

n = 0             # number of particles that have arrived but not decayed
eventList = []             #  See documentation.  The event list starts ...
heapq.heapify(eventList)   #  .. as a list, then is made into a heap.

for t_o in t_obs:                      #  schedule all the observations
   event = ( t_o, "observation")       #  t_o = time of observation
   heapq.heappush( eventList, event)
k_obs = 0                              # k_obs labels the observations
   
t = rng.exponential(alpha)     # generate the first arrival time
event = ( t, "arrival")        # the first arrival event
heapq.heappush( eventList, event)

while ( True):

   t, e_type = heapq.heappop(eventList)  # the time and type of next event
   if ( t > Tmax ):
      break
   
   if ( e_type == "arrival"):             # an arrival
      n  = n+1
      
      td = t + rng.exponential(beta)      # schedule this particle decay
      decay_event = ( td, "decay")        # create the decay event
      heapq.heappush( eventList, decay_event)  # put it in the event list
      
      ta = t + rng.exponential(alpha)         # schedule next arrival
      arrival_event = ( ta, "arrival")        # create the arrival event
      heapq.heappush( eventList, arrival_event)  # put it in the event list
      
   if ( e_type == "decay"):               # a decay
      n = n-1
      
   if ( e_type == "observation"):         # observation time, record n
      n_vals[k_obs] = n
      k_obs = k_obs+1
   
#   plot the n(t) process

fig,ax = plt.subplots()

ax.plot(t_obs, n_vals)

ax.grid()
titleStart = r"Arrival and decay process with "
titleEnd   = r"$\lambda =${lam:5.2f}, $\mu =${mu:5.2f}"
titleEnd = titleEnd.format(lam = lam, mu = mu)
titleString = titleStart + titleEnd
ax.set_title(titleString)
ax.set_ylabel("n(t)")
ax.set_xlabel("t")
plt.savefig("ArrivalDecayProcess.pdf")