# This is the R version of program 6.4 from page 203 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.

# This is a more complex stochastic model as 6 events are
# possible: infection, recovery, birth, death of susceptible, death of
# infected, death of recovered.

# 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,6,1)
  Change = matrix(0,6,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)
  
  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,]
  
  Iterate = c(step,new_value)
}

# Do the iterations using the full event driven stochastic methodology
# Note that this could be written more generally (see Program 6.4) but
# this version is kept simple for clarity

Stoch_Iteration = function(time,initial,pars){
  
  library(deSolve)
  
  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)
     step=Iterate.out[1]
     new=Iterate.out[2:4]
     
     loop=loop+1
     TT[loop]=TT[loop-1]+step
     X[loop] = old[1]
     Y[loop] = old[2]
     Z[loop] = old[3]
     
     loop=loop+1
     TT[loop]=TT[loop-1]
     X[loop] = new[1]
     Y[loop] = new[2]
     Z[loop] = new[3]
     
     old=new
     
   }
   
    Stoch_Iteration = cbind(TT,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

# Initial conditions
N0=5000      # Population size - assumed to be constant
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=2*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(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,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)

# 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,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))

p3=ggplot() +
  geom_line(data=out,aes(y=Z,x= TT/365),
            linewidth=0.3,linetype=1,color="black") +
  xlab("Time (years)") + ylab("Recovereds") +
  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))

library(patchwork)
p1/p2/p3