#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Dec 19 17:20:05 2022
%file stochastic_gene_expression.m
%stochastic implementation (Gillespie SSA) of a constitutive gene
%expression
%Figure 7.46

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



#set kinetic parameter values
kr=10;
kp=6;
gr=1;
gp=1;


#set initial condition for molecule counts (M, P) and 
#for simulation time

X=np.zeros((10000,3))

#set number of reactions
Tend=10000

for j in range(Tend-1):

    #calculate propensities
    a1=kr
    a2=kp*X[j,1]
    a3=gr*X[j,1]
    a4=gp*X[j,2]

    asum=a1+a2+a3+a4

    #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]+1
        X[j+1,2]=X[j,2]
    elif a1/asum <= mu and mu  < (a1+a2)/asum:
        X[j+1,1]=X[j,1]
        X[j+1,2]=X[j,2]+1
    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]=X[j,2]
    else:
        X[j+1,1]=X[j,1]
        X[j+1,2]=max(X[j,2]-1,0)
    

#produce figure 7.46A

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



ax1.plot(X[:,0], X[:,1], linewidth=2)
ax1.plot(X[:,0], X[:,2], linewidth=2)
    
ax1.set_xlim(0,65)
ax1.set_ylim(0,140)
ax1.set_ylabel('Number of Molecules')
ax1.set_xlabel('Time')



# %modify burst size
b=5

kr=10/b
kp=6
gr=1
gp=1

for j in range(Tend-1):

    #calculate propensities
    a1=kr
    a2=kp*X[j,1]
    a3=gr*X[j,1]
    a4=gp*X[j,2]

    asum=a1+a2+a3+a4

    #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]+5
        X[j+1,2]=X[j,2]
    elif a1/asum <= mu and mu  < (a1+a2)/asum:
        X[j+1,1]=X[j,1]
        X[j+1,2]=X[j,2]+1
    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]=X[j,2]
    else:
        X[j+1,1]=X[j,1]
        X[j+1,2]=max(X[j,2]-1,0)
    

#produce figure 7.46B



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