#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Dec 17 13:09:41 2022
%file stochastic_decay.m
%stochastic implementation (Gillespie SSA) of a collection of spontaneously
%decaying molecules.
%Figure 7.38

@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

#declare vectors for storing the steps in molecule count and time
X=np.zeros(1000)
t=np.zeros(1000)

#set parameter value
k=1

#set initial condition for molecule count and for simulation time.
X[0]=10
t[0]=0

#run SSA steps until 5 time units have passed
i=0
while t[i] < 5:
    #calculate propensity
    if X[i]>0:
        a=k*X[i]
    
    #update counter
    i=i+1

    #update time
    t[i]=t[i-1]+np.log(1/random.uniform(0, 1))/a

    #decrement state
    X[i]=max(X[i-1]-1,0)


#truncate zero values
#del X[i:]
#del t[i:]


#produce Figure 7.38C
fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
ax3.step(t,X, linewidth=2)
ax3.set_ylabel('Number of Molecules')
ax3.set_xlabel('Time')
ax3.set_xlim(0, 5)
ax3.set_ylim(0, 10)
#deterministic plot
det_t=np.linspace(0,5,100)
sol=10*np.exp(-det_t)
ax3.plot(det_t, sol)

#repeat for Figure 7.38B


#declare vectors for storing the steps in molecule count and time
X=np.zeros(1000)
t=np.zeros(1000)

#set parameter value
k=1

#set initial condition for molecule count and for simulation time.
X[0]=100
t[0]=0

#run SSA steps until 5 time units have passed
i=0
while t[i] < 5:
    #calculate propensity
    if X[i]>0:
        a=k*X[i]
    
    #update counter
    i=i+1

    #update time
    t[i]=t[i-1]+np.log(1/random.uniform(0, 1))/a

    #decrement state
    X[i]=max(X[i-1]-1,0)


#truncate zero values
#del X[i:]
#del t[i:]


#produce Figure 7.38C
ax2.step(t,X, linewidth=2)
ax2.set_ylabel('Number of Molecules')
ax2.set_xlabel('Time')
ax2.set_xlim(0, 5)
ax2.set_ylim(0, 100)
#deterministic plot
det_t=np.linspace(0,5,100)
sol=100*np.exp(-det_t)

ax2.plot(det_t, sol)



# %repeat for Figure 7.38C


#declare vectors for storing the steps in molecule count and time
X=np.zeros(1000)
t=np.zeros(1000)

#set parameter value
k=1

#set initial condition for molecule count and for simulation time.
X[0]=1000
t[0]=0

#run SSA steps until 5 time units have passed
i=0
while t[i] < 5:
    #calculate propensity
    if X[i]>0:
        a=k*X[i]
    
    #update counter
    i=i+1

    #update time
    t[i]=t[i-1]+np.log(1/random.uniform(0, 1))/a

    #decrement state
    X[i]=max(X[i-1]-1,0)


#truncate zero values
#del X[i:]
#del t[i:]


#produce Figure 7.38A
ax1.step(t,X, linewidth=2)
ax1.set_ylabel('Number of Molecules')
ax1.set_xlabel('Time')
ax1.set_xlim(0, 5)
ax1.set_ylim(0, 1000)
#deterministic plot
det_t=np.linspace(0,5,100)
sol=1000*np.exp(-det_t)

ax1.plot(det_t, sol)