#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Dec 19 17:20:05 2022
% file stochastic_vilar_oscillator.py
% stochastic simulation of genetic oscillator
% from Vilar et al. (2002) PNAS 99 5988-5992.
% problem 7.8.27


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



#set kinetic parameter values
eps=0.1
ga=250
ba=5
Ka=0.5 
alpha0=0.1
da=1
kc=200
gr=50
br=10
Kr=1
dr=eps*da


#set initial condition for molecule counts  and 
#for simulation time state: [time, A, R, C]

X=np.zeros((20000,4))

#set number of reactions
Tend=20000

for j in range(Tend-1):
    
    A=X[j,1]
    R=X[j,2] 
    C=X[j,3]

    #calculate propensities
    a1=(ga/ba)*(alpha0+(A/Ka))/(1+(A/Ka))
    a2=da*A
    a3=kc*A*R
    a4=(gr/br)*((A/Kr))/(1+(A/Kr))
    a5=dr*R
    a6=da*C
    
    asum=a1+a2+a3+a4+a5+a6

    #update time
    X[j+1,0]=X[j,0]+np.log(1/random.uniform(0,1))/asum

    #state transition
    mu=random.uniform(0,1)
    if 0 <= mu and mu < a1/asum:
        X[j+1,1]=X[j,1]+ba
        X[j+1,2]=X[j,2]
        X[j+1,3]=X[j,3]
    elif a1/asum <= mu and mu  < (a1+a2)/asum:
        X[j+1,1]=max(X[j,1]-1,0)
        X[j+1,2]=X[j,2]
        X[j+1,3]=X[j,3]
    elif (a1+a2)/asum <= mu and mu < (a1+a2+a3)/asum:
        X[j+1,1]=max(X[j,1]-1,0)
        X[j+1,2]=max(X[j,2]-1,0)
        X[j+1,3]=X[j,3]+1
    elif (a1+a2+a3)/asum <= mu and mu < (a1+a2+a3+a4)/asum:
        X[j+1,1]=X[j,1]
        X[j+1,2]=X[j,2]+br
        X[j+1,3]=X[j,3]
    elif (a1+a2+a3+a4)/asum <= mu and mu < (a1+a2+a3+a4+a5)/asum:
        X[j+1,1]=X[j,1]
        X[j+1,2]=max(X[j,2]-1,0)
        X[j+1,3]=X[j,3]
    else:
        X[j+1,1]=X[j,1]
        X[j+1,2]=X[j,2]+1
        X[j+1,3]=max(X[j,3]-1,0)

#produce figure 7.46A

fig, (ax1, ax2) = plt.subplots(1, 2)



ax1.plot(X[:,0], X[:,1], label="A", linewidth=2)
ax1.set_ylabel('Number of Molecules')
ax1.set_xlabel('Time')

ax2.plot(X[:,0], X[:,2], linewidth=2)
ax2.set_ylabel('Number of Molecules')
ax2.set_xlabel('Time')