# This is the R version of program 6.3 from page 202 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.

# It is the SIS model with full (event-driven) demographic stochasticity.

# This is a relatively simple stochastic model as only 2 events are
# possible: infection or recovery.
# Note: by default we are using a very small population size to highlight
# the stochasticity.

# Do the iteration step

Iterate = function(Y,pars){
  
  Rate_Infection = beta*(N0-Y)*Y/N0
  Rate_Recovery = gamma*Y
  
  R1 = runif(1)
  R2 = runif(1)
  
  step = -log(R2)/(Rate_Infection+Rate_Recovery)
  
  if(R1<Rate_Infection/(Rate_Infection+Rate_Recovery)){
    new_value=Y+1 # do infection
  }
  else{
    new_value = Y-1 # do recovery
  }
  
  Iterate = cbind(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 = c(0)
  Y = initial
  P = c(Y)
  loop = 1
  
  while (TT[loop]<time[2] & Y>0 ) {
    Iterate.out = Iterate(Y,pars)
    step=Iterate.out[1]
    new=Iterate.out[2]
    loop=loop+1
    TT[loop]=TT[loop-1]+step
    P[loop]=Y
    loop=loop+1
    TT[loop]=TT[loop-1]
    P[loop]=new
    Y=new
    
    a = length(TT)
    if(loop>=a){
      TT[loop*2]=0
      P[loop*2]=0
    }
  }
   TT = TT[1:loop]
   P = P[1:loop]
  
  
  Stoch_Iteration = cbind(TT,P)
}  


# Set parameters (all rates are per day)
beta=0.03              # Transmission rate - incorporates the encounter rate 
                       # between susceptible and infectious individuals together 
                       # with the probability of transmission
gamma=1.0/100.0        # Recovery rate

# Initial conditions
Y0=70      # Initial number of infectious individuals
N0=100      # Population size - assumed to be constant

# Time range
MaxTime=10*365

# Check all parameters are valid
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(MaxTime<=0){
  print(paste("Max run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if(Y0>N0){
  print(paste("Initial level of infecteds is greater than the population size, Y0 =",Y0,", N0 = ",N0))  
  stop()}


pars=c(beta,gamma,N0)

time = c(0, MaxTime)
initial = Y0

# 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=P,x=TT/365),
            linewidth=0.3,linetype=1,color="red") +
  xlab("Time (years)") + ylab("Susceptibles") +
  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