# This is the R version of program 5.2 from page 171 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the simple SIR epidemic with corrected term-time forcing of the transmission rate.
# Note: setting b1 too high can cause numerical difficulties.

# This code can also be used to generate bifurcation diagrams, by setting
# b1 equal to a vector of seasonality rates. The bifurcation diagram is
# constructed using extrapolated initial conditions. Try:
# (beta0,b1,gamma,mu,S0,I0,MaxYears) = (17/13,[0:0.01:0.25],1/13,1/(50*365),1/17,1e-4,20)

# Sets up term times
term = function(t){
  t = t%%365
  if((t<6) | (t>100 & t<115) | (t>200 & t<251) | (t>300 & t<307) | (t>356)){
    term = -1
  } else {
    term = 1
  }
} 

# Set up differential equations as a function
SIR_term_forcing = function(t,state,pars){
  with(as.list(c(state,pars)),{
    
    dSdt = mu - beta*S*I - mu*S
    dIdt = beta*S*I - gamma*I - mu*I
    dRdt = gamma*I - mu*R
    
    return(list(c(dSdt,dIdt,dRdt)))
  })
}

# Calculates the differential rates used in the integration
FORCED_ODE = function(state,pars){
  
  pars_minus = c(beta=(Beta0-B1),gamma,mu)
  pars_plus = c(beta=(Beta0+B1),gamma,mu)
  
  # Solve differential equations defined in function SIR_term_forcing
  library(deSolve)
  FORCED = c()
  
  for (Year in 0:(MaxYears-1)){
    times = seq(Year*365+0,Year*365+6,step)
    out1 = ode(state,times,SIR_term_forcing,pars_minus,rtol=1e-8)
    last_row = tail(out1,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+6,Year*365+100,step)
    out2 = ode(state,times,SIR_term_forcing,pars_plus,rtol=1e-8)
    last_row = tail(out2,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+100,Year*365+115,step)
    out3 = ode(state,times,SIR_term_forcing,pars_minus,rtol=1e-8)
    last_row = tail(out3,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+115,Year*365+200,step)
    out4 = ode(state,times,SIR_term_forcing,pars_plus,rtol=1e-8)
    last_row = tail(out4,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+200,Year*365+251,step)
    out5 = ode(state,times,SIR_term_forcing,pars_minus,rtol=1e-8)
    last_row = tail(out5,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+251,Year*365+300,step)
    out6 = ode(state,times,SIR_term_forcing,pars_plus,rtol=1e-8)
    last_row = tail(out6,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+300,Year*365+307,step)
    out7 = ode(state,times,SIR_term_forcing,pars_minus,rtol=1e-8)
    last_row = tail(out7,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+307,Year*365+356,step)
    out8 = ode(state,times,SIR_term_forcing,pars_plus,rtol=1e-8)
    last_row = tail(out8,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    times = seq(Year*365+356,Year*365+365,step)
    out9 = ode(state,times,SIR_term_forcing,pars_minus,rtol=1e-8)
    last_row = tail(out9,n=1)
    state = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    FORCED = rbind(FORCED,head(out1,-1),head(out2,-1),head(out3,-1),head(out4,-1),
          head(out5,-1),head(out6,-1),head(out7,-1),head(out8,-1),head(out9,-1))
    }
  FORCED_ODE = FORCED
}

# Set parameters (all rates are per day)
beta0=17/13      # Mean transmission rate
b1=0.25          # Amplitude of term-time forcing
# For bifurcation plot set b1 to the following vector
b1 = seq(from = 0, to = 0.25, by = 0.01)
gamma=1/13.0     # Recovery rate
mu=1/(50*365.0)  # Per capita deaty rate/population level birth rate

# Initial conditions
S0=1/17   # Initial proportion of population susceptible
I0=1e-4   # Initial proportion of population infectious
R0 = 1-S0-I0
y0 = c(S=S0,I=I0,R=R0)

# Time range
MaxYears=10
# For bifurcation plot set MaxYears to the following value
MaxYears=20
step=1

# Check all parameters are valid
if(S0<=0){
  print(paste("Initial level of susceptibles is less than or equal to zero, S0 =",S0))  
  stop()}
if(I0<=0){
  print(paste("Initial level of infecteds is less than or equal to zero, I0 =",I0))  
  stop()}
if(beta0<=0){
  print(paste("Transmission rate is less than or equal to zero, beta0 =",beta0))  
  stop()}
if(any(b1<0)==TRUE){
  print(paste("Seasonality b1 has an element less than zero"))  
  print(paste("b1 =")); print(b1)
  stop()}
if(any(b1>1)==TRUE){
  print(paste("Seasonality b1 has an element that is greater than or equal to 1"))  
  print(paste("b1 =")); print(b1)
  stop()}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(mu<0){
  print(paste("Birth/death rate is less than zero, mu =",mu))  
  stop()}
if(MaxYears<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxYears =",MaxYears))  
  stop()}
if((S0+I0)>1){
  print(paste("Initial level of susceptibles+infecteds is greater than 1, S0 =",S0,", I0 = ",I0))  
  stop()}
if(beta0<(gamma+mu)){
  print(paste("Basic reproductive ratio is less than 1, R_0 =",beta0/(gamma+mu)))  
  stop()}

# The main iteration
if(length(b1)==1){
  # Calculate average effect of forcing and correct for it
  Ave = 0
  for (t in 1:365){
    Ave=Ave+(1+b1*term(t-0.5))
  }
  Beta0 = beta0/(Ave/365)
  B1 = b1
  pars = c(Beta0,B1,gamma,mu,MaxYears,step) 

  FORCED.out = FORCED_ODE(y0,pars)
  out = as.data.frame(FORCED.out)

  # Plot the results using ggplot2
  library(ggplot2)
  p1=ggplot() +
    geom_line(data=out,aes(y=S,x= time/365),
              linewidth=0.6,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)),breaks=c(0,1,2,3,4,5,6,7,8,9,10)) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black", 
                                      fill = NA, linewidth = 0.6))
  p2=ggplot() +
    geom_line(data=out,aes(y=I,x= time/365),
              linewidth=0.6,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)),breaks=c(0,1,2,3,4,5,6,7,8,9,10)) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black", 
                                      fill = NA, linewidth = 0.6))
  
  p3=ggplot() +
    geom_line(data=out,aes(y=R,x= time/365),
              linewidth=0.6,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)),breaks=c(0,1,2,3,4,5,6,7,8,9,10)) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black", 
                                      fill = NA, linewidth = 0.6))
  
  library(patchwork)
  p1/p2/p3

} else {
  if(MaxYears*365<3650){
    MaxYears=10
  }  
  times=seq(0,MaxYears*365,step)
  Bifur_I = matrix(rep(0,length(b1)*10), 
                   nrow = length(b1), ncol = 10, byrow = TRUE) 
  Last = c(S=S0,I=I0,R=R0)
  
  for (loop in 1:length(b1)){
    B1=b1[loop]
    Ave=0
    for (t in 1:365){
      Ave=Ave+(1+B1*term(t-0.5))
    }
    Beta0=beta0/(Ave/365)
    pars = c(Beta0,B1,gamma,mu,MaxYears,step) 
    
    FORCED.out = FORCED_ODE(Last,pars)
    
    last_row = tail(FORCED.out,n=1)
    Last = c(S=last_row[2],I=last_row[3],R=last_row[4])
    
    infectious = FORCED.out[seq(1,MaxYears*365,365),3]
    
    # Just save last 10 years
    Bifur_I[loop,1:10] = infectious[(MaxYears-9):MaxYears]
  }
  library(tidyr)
  out = data.frame(b1,Bifur_I)
  out_long=pivot_longer(out,cols=c("X1","X2","X3","X4","X5","X6","X7","X8","X9","X10"),
                        names_to="Bifur",values_to="infectious")
  
  # Plot the results using ggplot2
  library(ggplot2)
  
  p1=ggplot() +
    geom_point(data=out_long,aes(y=infectious,x=b1),
               size=0.4,color="black") +
    xlab(bquote("Seasonality, "~b[1])) +
    ylab("Level of infection") +
    scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
    scale_x_continuous(expand=expansion(mult=c(0.0,0.01)),limits=c(-0.001,0.251)) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black", 
                                      fill = NA, linewidth = 0.6))
  
  library(patchwork)
  p1
  
}