# This is the R version of program 5.3 from page 184 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the simple SIR epidemic with sinusoidal forcing of the birth rate.
# Note: setting alpha1 too high can cause numerical difficulties.

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

# Set up differential equations as a function
SIR_sin_births = function(t,state,pars){
  with(as.list(c(state,pars)),{
    
    alpha = alpha0*(1+alpha2*sin(2.0*pi*t/365.0))
    mu = alpha0
    
    dSdt = alpha - beta*S*I - mu*S
    dIdt = beta*S*I - gamma*I - mu*I
    dRdt = gamma*I - mu*R
    
    return(list(c(dSdt,dIdt,dRdt)))
  })
}


# Set parameters (all rates are per day)
beta=17/13          # Transmission rate
gamma=1/13.0        # Recovery rate
alpha0=1/(50*365.0) # Mean birth rate
alpha1=0.25         # Amplitude of sinusoidal forcing for birth rate
# For bifurcation plot set alpha1 to the following vector
#alpha1 = seq(from = 0, to = 0.99, by = 0.01)

# 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
MaxTime=60*365
# For bifurcation plot set MaxTime to the following value
#MaxTime=20*365
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(beta<=0){
  print(paste("Transmission rate is less than or equal to zero, beta =",beta))  
  stop()}
if(any(alpha1<0)==TRUE){
  print(paste("Seasonality alpha1 has an element less than zero"))  
  print(paste("alpha1 =")); print(alpha1)
  stop()}
if(any(alpha1>1)==TRUE){
  print(paste("Seasonality alpha1 has an element that is greater than or equal to 1"))  
  print(paste("alpha1 =")); print(alpha1)
  stop()}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(alpha0<0){
  print(paste("Birth/death rate is less than zero, alpha0 =",alpha0))  
  stop()}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if((S0+I0)>1){
  print(paste("Initial level of susceptibles+infecteds is greater than 1, S0 =",S0,", I0 = ",I0))  
  stop()}

# The main iteration
if(length(alpha1)==1){
  alpha2 = alpha1
  pars = c(beta,gamma,alpha0,alpha2)
  times = seq(0,MaxTime,step)
  
  # Solve differential equations defined in function SIR_sin_forcing
  library(deSolve)
  SIR_sin_births.out = ode(y0,times,SIR_sin_births,pars,rtol=1e-5)
  out = as.data.frame(SIR_sin_births.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))) +
    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))) +
    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))) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black", 
                                      fill = NA, linewidth = 0.6))
  
  library(patchwork)
  p1/p2/p3
  
} else {
  if(MaxTime<3650){
    MaxTime = 3650
  }
  times =seq(0,MaxTime,365)
  Bifur_I = matrix(rep(0,length(alpha1)*10), 
                   nrow = length(alpha1), ncol = 10, byrow = TRUE) 
  
  Last = c(S=S0,I=I0,R=R0)
  for (loop in 1:length(alpha1)){
    alpha2 = alpha1[loop]
    pars = c(beta,gamma,alpha0,alpha2)
    
    # Solve differential equations defined in function SIR_sin_forcing
    library(deSolve)
    SIR_sin_births.out = ode(Last,times,SIR_sin_births,pars,rtol=1e-5)
    
    last_row = tail(SIR_sin_births.out,n=1)
    Last = c(S=last_row[2],I=last_row[3],R=last_row[4])
    infectious = SIR_sin_births.out[,3]
    # Just save last 10 years
    Bifur_I[loop,1:10] = infectious[(length(times)-9):length(times)]
  }   
  library(tidyr)
  out = data.frame(alpha1,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= alpha1),
               size=0.5,color="black") +
    xlab(bquote("Seasonality, "~alpha[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.01))) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black", 
                                      fill = NA, linewidth = 0.6))
  
  library(patchwork)
  p1
  
}