# This is the R version of program 2.1 from page 19 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the simple SIR epidemic without births or deaths.
 

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

# Set parameters (all rates are per day)
beta=1.4247       # Transmission rate
gamma=0.1426      # Recovery rate
pars = c(beta,gamma)

# Initial conditions
S0 = 1.0-1e-6  # initial proportion of population that are susceptible
I0 = 1e-6      # initial proportion of population that are infectious
y0 = c(S=S0,I=I0,R=0)

# Time range
MaxTime=70
step = 0.1
times = seq(0,MaxTime,step)

# Check all parameters are valid
if(S0<=0){
  print(paste("Initial level of susceptibles is less than or equal to zero, S0 =",S0))  
}
if(I0<=0){
  print(paste("Initial level of infecteds is less than or equal to zero, I0 =",I0))  
}
if(beta<=0){
  print(paste("Transmission rate is less than or equal to zero, beta =",beta))  
}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
}
if(S0+I0>1){
  print(paste("Initial level of S+I is greater than 1, S0+I0  =",S0+I0))  
}
if(beta<gamma){
  print(paste("Basic reproductive ratio is less than 1, R =",beta/gamma))
}

# Solve differential equations
library(deSolve)
SIR.out = ode(y0,times,SIR, pars)

# Plot results using ggplot2
library(ggplot2)
out = as.data.frame(SIR.out)

p1=ggplot() +
  geom_line(data=out,aes(y=S,x= time,color="Y1"),linewidth=0.6,linetype=1) +
  geom_line(data=out,aes(y=R,x= time,color="Y2"),linewidth=0.6,linetype=1) +
  scale_color_manual(name = "", labels=c("S","R"),
                     values = c("Y1" = "green3", "Y2" = "black")) +
  xlab("Time") + ylab("Susceptibles and Recovereds") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05))) +
  scale_x_continuous(limits=c(0,70),expand=expansion(mult=c(0.0,0.0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.9,0.7),
       legend.box.background = element_rect(color = "black"),
       legend.spacing.y = unit(0, "mm")) +
       theme(panel.border = element_rect(color = "black", 
                      fill = NA, linewidth = 0.6))

p2=ggplot() +
  geom_line(data=out,aes(x=time,y=I,color="red"),linewidth=0.6,linetype=1) +
  xlab("Time") + ylab("Infectious") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05))) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(legend.position="none") +
  theme(panel.border = element_rect(color = "black", 
                 fill = NA, linewidth = 0.6))

library(patchwork)
p1/p2