#   This is the R version of program 5.1 from page 160 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
#  
# It is the simple SIR epidemic with sinusoidal forcing of the transmission rate.
# Note: setting beta1 too high can cause numerical difficulties.
# 
# This code can also be used to generate bifurcation diagrams, by setting
# beta1 equal to a vector of seasonality rates. The bifurcation diagram is
# constructed using extrapolated initial conditions. Try:
#  (beta0,beta1,gamma,mu,S0,I0,MaxTime) =
#      (17/13,[0:0.001:0.25],1/13,1/(50*365),1/17,1e-4,20*365);

# Set up differential equations as a function
SIR_sin_forcing = function(t,state,pars){
  with(as.list(c(state,pars)),{
    
    beta = beta0*(1+beta2*sin(2.0*pi*t/365.0))
    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)))
  })
}

# Set parameters (all rates are per day)
beta0=17/13     # Mean transmission rate
beta1=0.1       # Amplitude of sinusoidal forcing
# For bifurcation plot set beta1 to the following vector
beta1 = seq(from = 0, to = 0.25, by = 0.001)
gamma=1/13.0    # Recovery rate
mu=1/(50*365.0) # Per capita death 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
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(beta0<=0){
  print(paste("Transmission rate beta0 is less than or equal to zero, beta0 =",beta0))  
  stop()}
if(any(beta1<0)==TRUE){
  print(paste("Seasonality beta1 has an element less than zero"))  
  print(paste("beta1 =")); print(beta1)
  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(S0+I0>1){
  print(paste("Initial level of susceptibles+infecteds is greater than 1, S0+I0 =",S0+I0))  
  stop()}
if(beta0<(gamma+mu)){
  print(paste("Basic reproductive ratio is less than 1, beta0/(gamma+mu) =",beta0/(gamma+mu)))  
  stop()}

# The main iteration
if(length(beta1)==1){
  beta2 = beta1
  pars = c(beta0,beta2,gamma,mu)
  times = seq(0,MaxTime,step)
  # Solve differential equations defined in function SIR_sin_forcing
  library(deSolve)
  SIR_sin_forcing.out = ode(y0,times,SIR_sin_forcing,pars,rtol=1e-5)
  out = as.data.frame(SIR_sin_forcing.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(beta1)*10), 
                   nrow = length(beta1), ncol = 10, byrow = TRUE) 
  
  Last = c(S=S0,I=I0,R=R0)
  for (loop in 1:length(beta1)){
    beta2 = beta1[loop]
    pars = c(beta0,beta2,gamma,mu)
    
  # Solve differential equations defined in function SIR_sin_forcing
    library(deSolve)
    SIR_sin_forcing.out = ode(Last,times,SIR_sin_forcing,pars,rtol=1e-5)
  
    last_row = tail(SIR_sin_forcing.out,n=1)
    Last = c(S=last_row[2],I=last_row[3],R=last_row[4])
    infectious = SIR_sin_forcing.out[,3]
    # Just save last 10 years
    Bifur_I[loop,1:10] = infectious[(length(times)-9):length(times)]
  }   
    library(tidyr)
    out = data.frame(beta1,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= beta1),
                size=0.5,color="black") +
      xlab(bquote("Seasonality, "~beta[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
  
}