#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Dec 19 16:15:58 2022
%file stochastic_decay_ensemble.m
%stochastic implementation (Gillespie SSA) of a collection of spontaneously
%decaying molecules.
%Figure 7.45

@author: bingalls
"""
import numpy as np
import matplotlib.pyplot as plt
import random



#the network is simply  X -> . with propensity a=k N_x, 
#where N_x is the number of molecules of X



#set parameter value
k=1

#set ensemble size
N=3

atot=1

#set initial condition for molecule count and for simulation time, 
#for each member of ensemble
X1=np.zeros((10,N))
t1=np.zeros((10,N))
X1[0,:]=10
t1[0,:]=0

#generate ensemble
for j in range(N):

    #run ten steps of the simulation algorithm
    for i in range(10-1):
    
        #calculate propensity
        if X1[i,j]>0:
            atot=k*X1[i,j]

        #update time
        t1[i+1,j]=t1[i,j]+np.log(1/random.uniform(0,1))/atot

        #decrement state
        X1[i+1,j]=max(X1[i,j]-1,0)


#collect vector of reaction time events 
rxntimes=[]
for j in range(N):
    rxntimes.append(t1[:,j])


#interpolate molecule counts for plotting 
# #(otherwise the ensemble of staircase graphs lie on top of one another)
# from scipy.interpolate import interp1d
# for j in range(N):
#     X1interp[i,j]=interp1d(t1[:,j], X1[:,j], rxntimes);



# #%produce Figure 7.45A

fig, (ax1, ax2, ax3) = plt.subplots(1, 3)



#ax1.plot(rxntimes, X1mean, linewidth=2)
for j in range(N):
    ax1.plot(t1[:,j], X1[:,j], linewidth=1)
    
ax1.set_xlim(0,5)
ax1.set_ylim(0,10)
ax1.set_ylabel('Number of Molecules')
ax1.set_xlabel('Time')

# %repeat for larger ensemble sizes