#      This is the R version of program 6.5 from page 204 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SIR model (including births and deaths) with 
# (event-driven) demographic stochasticity approximated using the tau-leap 
# method and assuming Poisson distributions.

# 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
Tau_leap = function(old,pars){
  
  X=old[1]
  Y=old[2]
  Z=old[3]
  
  Rate = matrix(0,6,1)
  Change = matrix(0,6,3)
  
  Rate[1] = mu*N0
  Change[1,]=c(1,0,0)
  Rate[2] = beta*X*Y/N0
  Change[2,]=c(-1,1,0)
  Rate[3] = gamma*Y
  Change[3,]=c(0,-1,1)
  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)
  
  new_value=old

    for (i in 1:6){
    Num=rpois(1,Rate[i]*tau)
    # Make sure things don't go negative
    Use=min(c(Num,new_value[which(Change[i,]<0)]))
    new_value=new_value+Change[i,]*Use
  }
  Tau_leap = 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)
  
  TT = matrix(0,time[2]+1,1)
  P = matrix(0,time[2]+1,3)
  colnames(P) = c('X','Y','Z')
  
  TT=seq(0,time[2],tau)
  P[1,]=initial
  
  X = initial[1]
  Y = initial[2]
  Z = initial[3]
  old=c(X,Y,Z)
  
  loop = 1
  while (TT[loop]<time[2]){
     Tau_leap.out = Tau_leap(old,pars)
     #step=Iterate.out[1]
     new=Tau_leap.out[1:3]
     
     loop=loop+1
     P[loop,]=new
     old=new
  }
    Stoch_Iteration = cbind(TT,P)
}  


# 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

# 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=2*365
tau = 1             # Time step

# 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(tau<=0){
  print(paste("Step length tau is less than or equal to zero, tau =",tau))  
  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,tau)

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