#  Code for Stochastic Calculus class, Jonathan Goodman
#  http://www.math.nyu.edu/faculty/goodman/teaching/StochCalc2021/
#  ExponentialSampler.py

#  Demo of random sampling, histograms, and programming practice
#  See Exercise 5 of Week 1 notes for details

# illustrate using the Python random number generator from numpy version 19

from   numpy.random import Generator, PCG64       #  numpy randon number generator routines
import numpy             as np
import matplotlib.pyplot as plt

#-------------------- The simulator, an exponential larger than ts --------------------

def sim( ts, lam, bg):
    """Simulate independent exponential random variables with rate lam
    until you get one larger than ts (t-start, starting time) and return
    that.  Arguments
    ts:  (type = float) return the first T ~ exp(lam) with T>ts
    lam: (type = float) the rate parameter for the exponential random variable
    bg:  a instance of the random number generator
    return: T, n, where T = the first exponential above ts, n = number of trials"""
    
    N = 1             #  Count number of trials
    while(1):         #  could be an infinite loop, dangerous
        T = -(1/lam)*np.log( rg.random() )  # Generate T ~ exp(lam),
                                            # See Week 1 Exercise 5 for this formula
        if ( T > ts ):                      # return it if it's above threshhold
            return T, N
        N += 1                              # count the T you just rejected
        
        
#------------------ The main program ---------------------------------------------------

bg = PCG64(12345)        # instantiate a bit generator with seed 12345
rg = Generator(bg)       # instantiate a random number generator

#    Collect all constants and parameters in one place

lam      = .136                       # the exponential rate constant
ts       = 2.48/lam                   # the threshhold
nhist    = 20000                      # number of samples for the histogram
maxT     = 5./lam                     # make bins out to 5*expected T
nbins    = 30                         # number of histogram bins
PlotFile = "ExponentialHistogram.pdf" # filename for histogram plot file

Tlist = np.zeros(nhist)               # allocate a numpy array for the histogram samples
Nsum  = 0                             # To calculate the average of N
for k in range(nhist):
    T, N     = sim( ts, lam, bg)
    Tlist[k] = T - ts                 # time to hit after ts
    Nsum    += N
    
DelT = maxT/nbins                     # bin width
binT = np.zeros(nbins)                # bin centers,
binC = np.zeros(nbins, dtype = int)   # bin counts are integers
for k in range(nhist):
    Tk = Tlist[k]
    if ( Tk < maxT):                  # count samples in the range of the histogram
        Bk = int( Tk/DelT )           # truncate to identify the bin
        binC[Bk] += 1

pe = np.zeros(nbins)             #  empirical probability estimate from simulation
pt = np.zeros(nbins)             #  theoretical probability density

Nbar = Nsum / nhist              #  average number of trials

infoLine = "\n Exponentials larger than {ts:10.3e} with rate {lam:10.3e}"
infoLine = infoLine.format( ts = ts, lam = lam)
print(infoLine)
infoLine = " Expected number of trials was {Nbar:10.3e} \n"
infoLine = infoLine.format( Nbar = Nbar)
print(infoLine)
print(" bin center  empirical probability  theoretical probability    difference\n")

for j in range(nbins):
    pe[j]   = binC[j]/(nhist*DelT)     #  See week 1 notes for this formula
    Tj      = ( j + .5)*DelT           #  The center of the bin
    binT[j] = Tj                       #  record the bin center times, for plotting
    pt[j]   = lam*np.exp(-lam*Tj)      #  exponential probability density
    
#      character string with formatting information
    
    table_line = "  {Tj:8.4f}         {pe:8.3e}               {pt:8.3e}         {diff:11.3e}"
    
#      put the numbers into the formatted string
    
    table_line = table_line.format( Tj = Tj, pe = pe[j], pt = pt[j], diff = pe[j]-pt[j])
    
#      print the line of the table, character string with formatted numbers embedded

    print(table_line)
        
fig, ax = plt.subplots()                                 # Create a figure containing a single axes.
ax.plot(binT, pe, 'bo', label = 'empirical density')     # add a curve with the empirical density estimates
ax.plot(binT, pt, label = 'theoretical density')         # add a curve with the theoretical densities
ax.legend()                                              # the legend says which graph is which
ax.set_ylabel('p')                                       # usually it's more complicated than this
ax.set_xlabel('T - ts')
title = "exponential rate {lam:6.2f}, {nf:8.2e} samples" # plot title with formatted numbers embedded
title = title + ", threshhold {ts:8.2e}"
title = title.format(lam = lam, nf = float(nhist), ts = ts)
ax.set_title(title)
ax.grid()                                                # the grid makes it easier to read the graph
plt.savefig(PlotFile)                                    # save in the same directory
#plt.show()