# This is the R version of program 2.6 from page 41 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SEIR epidemic with equal births and deaths.
# Note we no-longer explicitly model the recovered class.

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

# Set parameters (all rates are per day)
beta=520/365.0    # Transmission rate 
sigma = 1/14.0    # Rate at which individuals move from exposed to infectious class
gamma=1/7.0       # Recovery rate
mu=1/(70*365.0)   # Per capita death rate
pars = c(beta,sigma,gamma,mu)

# Initial conditions
S0 = 0.1       # initial proportion of population that are susceptible 
E0 = 1e-4      # initial proportion of population exposed (infected but not infectious)
I0 = 1e-4      # initial proportion of population that are infectious
y0 = c(S=S0,E=E0,I=I0)

# Time range
MaxTime=60*365
step=28
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(E0<=0){
  print(paste("Initial level of exposed is less than or equal to zero, E0 =",E0))  
}
if(I0<=0){
  print(paste("Initial level of infectious 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(sigma<=0){
  print(paste("Exposed to infectious rate is less than or equal to zero, sigma =",sigma))  
}
if(mu<0){
  print(paste("Birth/death rate is less than zero, mu =",mu))  
}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
}
if(S0+E0+I0>1){
  print(paste("Initial level of susceptibles+infecteds is greater than 1, S0+E0+I0  =",S0+E0+I0, "N0 =",N0))  
}
if(beta*sigma<(gamma+mu)*(sigma+mu)){
  print(paste("Basic reproductive ratio is less than 1, R =",beta*sigma/((gamma+mu)*(sigma+mu))))
}

# Solve the ODEs
library(deSolve)
SEIR.out = as.data.frame(ode(y0,times,SEIR,pars,rtol=1e-5))

# Plot the results using ggplot2
library(ggplot2)

p1=ggplot() +
  geom_line(data=SEIR.out,aes(y=S,x= time/1e4),linetype=1,color="green3") +
  xlab(bquote("Time " (x10^4))) + 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=SEIR.out,aes(y=E/1e-3,x= time/1e4,color="Y1"),linetype=1) +
  geom_line(data=SEIR.out,aes(y=I/1e-3,x= time/1e4,color="Y2"),linetype=1) +
  scale_color_manual(name = "", labels=c("Exposed","Infectious"),
                     values = c("Y1" = "red", "Y2" = "blue")) +
  xlab(bquote("Time " (x10^4))) + ylab(bquote("Infected " (x10^-3))) +
  scale_y_continuous(limits=c(0,1.2),expand=expansion(mult=c(0,0.05))) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.9,0.72),
        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))

p3=ggplot() +
  geom_line(data=SEIR.out,aes(y=1-S-E-I,x= time/1e4),
            linetype=1,color="black") +
  xlab(bquote("Time " (x10^4))) + 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