# This is the R version of program 2.7 from page 44 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SICR which includes a carrier class.

# Set up SEIR equations as a function
SICR = function(time,state,pars){
  with(as.list(c(state,pars)),{
    dSdt =  mu - beta*S*(I + epsilon*C) - mu*S
    dIdt = beta*S*(I + epsilon*C) - gamma*I- mu*I
    dCdt = gamma*q*I - Gamma*C - mu*C
    return(list(c(dSdt,dIdt,dCdt)))
  })
}

# Set parameters (all rates are per day)
beta=0.2          # Transmission rate 
epsilon=0.1       # proportion reduction in transmission from carriers compared to standard infectious
gamma=0.01        # Recovery rate
Gamma=0.001       # Recovery rate associated with carriers
mu=1/(50*365.0)   # Per capita death rate and population level birth rate
q=0.4             # Proportion of infected individuals that become carriers
pars = c(beta,epsilon,gamma,Gamma,mu,q)

# Initial conditions
S0 = 0.1       # initial proportion of population that are susceptible 
I0 = 1e-4      # initial proportion of population that are infectious
C0 = 1e-3      # initial proportion of population that are carriers
y0 = c(S=S0,I=I0,C=C0)

# 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(I0<=0){
  print(paste("Initial level of infectious is less than or equal to zero, I0 =",I0))  
}
if(C0<=0){
  print(paste("Initial level of carriers is less than or equal to zero, C0 =",C0))  
}
if(beta<=0){
  print(paste("Transmission rate is less than or equal to zero, beta =",beta))  
}
if(epsilon<=0){
  print(paste("Reduced transmission rate is less than or equal to zero, epsilon =",epsilon))  
}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
}
if(Gamma<=0){
  print(paste("Recovery rate of carriers is less than or equal to zero, Gamma =",Gamma))  
}
if(q<0 | q>1){
  print(paste("Probability of becoming a carrier is not a valid probability, q =",q))  
}
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+I0+C0>1){
  print(paste("Initial level of susceptibles+infecteds+carriers is greater than 1, S0+I0+C0  =",S0+I0+C0, "N0 =",N0))  
}
if(beta/(gamma+mu) + beta*q*epsilon/(Gamma+mu) < 1){
  print(paste("Basic reproductive ratio is less than 1, R =",beta/(gamma+mu) + beta*q*epsilon/(Gamma+mu)))
}

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

# Plot the results using ggplot2
library(ggplot2)

p1=ggplot() +
  geom_line(data=out,aes(y=S,x= time/1e4),linetype=1,color="green3") +
  xlab(bquote("Time " (x10^4))) + ylab("Susceptibles") +
  scale_y_continuous(limits=c(0,0.12),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/1e4),linetype=1,color="red") +
  xlab(bquote("Time " (x10^4))) + ylab("Infectious") +
  scale_y_continuous(limits=c(0,0.03),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=C,x= time/1e4),linetype=1,color="blue") +
  xlab(bquote("Time " (x10^4))) + ylab("Carriers") +
  scale_y_continuous(limits=c(0,0.04),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