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

# It is the SIR epidemic model with scaled additive noise added to 
# all the various rates.

# Set up differential equations as a function
SIR_scaled_noise = function(t,state,pars){
  with(as.list(c(state,pars)),{
    Noise=pars[2:6]
    dXdt = (mu*N0 + sqrt(mu*N0)*Noise[1]) - (beta*X*Y/N0 + sqrt(beta*X*Y/N0)*Noise[2]) - (mu*Y + sqrt(mu*Y)*Noise[3])
    dYdt = (beta*X*Y/N0 + sqrt(beta*X*Y/N0)*Noise[2]) - (gamma*Y + sqrt(gamma*Y)*Noise[4]) - (mu*Y + sqrt(mu*Y)*Noise[5])
    
    return(list(c(dXdt,dYdt)))
  })
}

# Do the integration of the differential equations with noise
# Here we fix the noise for each step, but integrate that step using ode

Stoch_ODE = function(time,initial,pars){
  library(deSolve)
  sqrtstep = sqrt(step)
  t = 0
  X = initial[1]
  Y = initial[2]
  
  TT = matrix(0,1+ceiling(time[2]/step),1)
  P = matrix(0,1+ceiling(time[2]/step),2)
  colnames(P) = c('X','Y')
  Noise = rep(0,5)
  
  P[1,] = initial
  
  loop = 1
  
  while (t<time[2] & Y>0 & X>0) {
    Noise[1:5] = rnorm(5)/sqrtstep
    times = seq(t,t+step,1)
    
    P0 = P[loop,]
    
    pars2=c(beta,Noise,gamma,mu,step,N0)
  
    out = ode(P0,times,SIR_scaled_noise,pars2,rtol=1e-5)
    last_row = tail(out,n=1)
    
    t = last_row[1]
    X = last_row[2]
    Y = last_row[3]
    
    loop = loop + 1
    TT[loop] = t
    P[loop,] = c(X=last_row[2],Y=last_row[3])
    
  }
  TT = TT[1:loop]
  P = P[1:loop,]
  
  
  Stoch_ODE = cbind(TT,P)
}  


# Set parameters (all rates are per day)
beta=1.0              # Transmission rate - incorporates encounter rate between susceptible and infectious individuals together with probability of transmission
gamma=1.0/10.0        # Recovery rate
mu=1.0/(50.0*365.0)   # Per capita death rate

# Initial conditions
X0=1e5      # Initial number of susceptible individuals
Y0=500      # Initial number of infectious individuals
N0=1e6      # Population size (constant)
P0 = c(X=X0,Y=Y0)

# Time range
MaxTime=5*365
step = 1

# 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/Death rate mu is less than zero, mu =",mu))  
  stop()}
if(MaxTime<=0){
  print(paste("Max run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if(step<=0){
  print(paste("Integration time-step is less than or equal to zero, step =",step))  
  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,step,N0)

time = c(0, MaxTime)
initial = P0

Stoch.out = Stoch_ODE(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="green3") +
  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))

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

library(patchwork)
p1/p2