#      This is the R version of program 6.6 from page 210 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SIR model (including births and deaths) with full 
# (event-driven) demographic stochasticity. Two forms of imports (governed
# by the parameters epsilon and delta) are implemented.

# This is a more complex stochastic model as 8 events are
# possible: infection, recovery, birth, death of susceptible, death of
# infected, death of recovered and two forms of imports

# Note: by default we are using a very small population size to highlight
# the stochasticity.

 # Do the iteration step
Iterate = function(old,pars){
  
  X=old[1]
  Y=old[2]
  Z=old[3]
  
  Rate = matrix(0,8,1)
  Change = matrix(0,8,3)
  
  Rate[1] = beta*X*Y/N0
  Change[1,]=c(-1,1,0)
  Rate[2] = gamma*Y
  Change[2,]=c(0,-1,1)
  Rate[3] = mu*N0
  Change[3,]=c(1,0,0)
  Rate[4] = mu*X
  Change[4,]=c(-1,0,0)
  Rate[5] = mu*Y
  Change[5,]=c(0,-1,0)
  Rate[6] = mu*Z
  Change[6,]=c(0,0,-1)
  Rate[7] = epsilon*X
  Change[7,] = c(-1,1,0)
  Rate[8] = delta
  Change[8,] = c(0,1,0)
  
  R1 = runif(1)
  R2 = runif(1)
  
  step = -log(R2)/(sum(Rate))
  
  # find which event to do
  m=min(which(cumsum(Rate)>=R1*sum(Rate)))
  new_value=old+Change[m,]
  
  dt = 0
  et = 0
  
  if(m==7){et = 1}
  if(m==8){dt = 1}
  
  Iterate = c(step,et,dt,new_value)

}

# Do the iterations using the full event driven stochastic methodology
# This is a relatively general version of the event-driven methodology
# known as Gillespie's Direct Algorithm

Stoch_Iteration = function(time,initial,pars){
  
  library(deSolve)
  
  ET = 0
  DT = 0
  TT = 0
  
  X = initial[1]
  Y = initial[2]
  Z = initial[3]
  old=c(X,Y,Z)
  
  loop = 1
  while (TT[loop]<time[2]){
    Iterate.out = Iterate(old,pars)
    loop=loop+1
    step = Iterate.out[1]
    TT[loop] = TT[loop-1]+step
    X[loop] = old[1]
    Y[loop] = old[2]
    Z[loop] = old[3]
    ET[loop] = Iterate.out[2]
    DT[loop] = Iterate.out[3]
    
    loop = loop + 1
    TT[loop] = TT[loop-1]
    new=Iterate.out[4:6]
    X[loop] = new[1]
    Y[loop] = new[2]
    Z[loop] = new[3]
    ET[loop] = Iterate.out[2]
    DT[loop] = Iterate.out[3]
    
    old=new
    
  }
  
  Stoch_Iteration = cbind(TT,ET,DT,X,Y,Z)
                            
}  


# Set parameters (all rates are per day)
beta=1            # Transmission rate - incorporates the encounter rate 
                       # between susceptible and infectious individuals together 
                       # with the probability of transmission
gamma=1.0/10.0    # Recovery rate
mu=5e-4           # Per capita death rate
N0 = 5000         # Population size - assumed to be constant
delta = 0.01      # Import rate due to external force of infection
epsilon = delta*10/N0  # Import rate due to infectious individuals immigrating into population

# Initial conditions
Y0=ceiling(mu*N0/gamma)    # Initial number of infectious individuals
X0=floor(gamma*N0/beta)    # Initial number of susceptible individuals
Z0 = N0-X0-Y0

# Time range
MaxTime=10*365

# Check all parameters are valid
if(X0<=0){
  print(paste("Initial number of susceptibles is less than or equal to zero, X0 =",X0))  
  stop()}
if(Y0<=0){
  print(paste("Initial number of infecteds is less than or equal to zero, Y0 =",Y0))  
  stop()}
if(N0<=0){
  print(paste("Initial population size is less than or equal to zero, N0 =",N0))  
  stop()}
if(beta<=0){
  print(paste("Transmission rate beta is less than or equal to zero, beta =",beta))  
  stop()}
if(gamma<=0){
  print(paste("Recovery rate gamma is less than or equal to zero, gamma =",gamma))  
  stop()}
if(mu<=0){
  print(paste("Birth and death rate mu is less than or equal to zero, mu =",mu))  
  stop()}
if(epsilon<=0){
  print(paste("Import rate epsilon is less than or equal to zero, epsilon =",epsilon))  
  stop()}
if(delta<=0){
  print(paste("Import rate delta is less than or equal to zero, delta =",delta))  
  stop()}
if(MaxTime<=0){
  print(paste("Max run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if(X0+Y0>N0){
  print(paste("Initial level of susceptibles+infecteds is greater than the population size, X0+Y0 =",X0+Y0,", N0 = ",N0))  
  stop()}


pars=c(beta,gamma,mu,epsilon,delta,N0)

time = c(0, MaxTime)
P0 = c(X=X0,Y=Y0,Z=Z0)
initial = P0

# The main iteration
Stoch.out = Stoch_Iteration(time,initial,pars)
out = as.data.frame(Stoch.out)
ET = out[,2]
DT = out[,3]

# Plot the results using ggplot2
library(ggplot2)

p1=ggplot() +
  geom_line(data=out,aes(y=X,x=TT/365),
            linewidth=0.3,linetype=1,color="green") +
  xlab("Time (years)") + ylab("Susceptible") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05))) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))


p2=ggplot() +
  geom_line(data=out,aes(y=Y,x= TT/365),
            linewidth=0.3,linetype=1,color="red") +
  xlab("Time (years)") + ylab("Infectious") +
 scale_y_continuous(expand=expansion(mult=c(0.05,0.05))) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))
Q = layer_scales(p2)$y$get_limits()
m = which(ET>0)
r = rep(0,length(m))
out2 = as.data.frame(cbind(r,out[m,1]))
n = which(DT>0)
s = rep(0,length(n))
s = s+Q[2]
out3 = as.data.frame(cbind(s,out[n,1]))
p2 = p2 + 
  geom_point(data=out2,aes(y=r,x=out2[,2]/365,fill="Epsilon Imports"),shape=21)+
  geom_point(data=out3,aes(y=s,x=out3[,2]/365,fill="Delta Imports"),shape=21) +
  scale_fill_manual(name = "", values = c("Epsilon Imports" = "blue", "Delta Imports" = "red")) +
  theme(legend.title = element_blank(),legend.position = c(0.87,0.7),
        legend.box.background = element_rect(color = "black"),
        legend.spacing.y = unit(0, "mm")) +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))
  
  
library(patchwork)
p1/p2