# This is the R version of program 3.5 from page 94 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the S(E)IR model with multiple stages to create gamma-distributed
# exposed and infectious periods
 
# n is the total number of infection classes; classes 1 to m are exposed,
# classes m+1 to n are infectious.
# Birth and death rates are assumed equal.
 
# As well as the default, you may want to compare the structured model:
# ( n, m, beta, gamma, mu, S0, I0, MaxTime ) = (10, 0, 1, 0.1, 0, 0.5, 1e-6, 60);
# with the unstructured version
#( n, m, beta, gamma, mu, S0, I0, MaxTime ) = (1, 0, 1, 0.1, 0, 0.5, 1e-6, 60);
# Or compare the SEIR:
# ( n, m, beta, gamma, mu, S0, I0, MaxTime ) = (10, 5, 1, 0.1, 0, 0.5, 1e-4, 150);
# with the unstructured version
# ( n, m, beta, gamma, mu, S0, I0, MaxTime ) = (2, 1, 1, 0.1, 0, 0.5, 1e-4, 150);
 
# Set up equations as a function
SEIR_mult = function(time,state,pars){
  with(as.list(c(state,pars)),{
    I = rep(0,n)
    dIdt = rep(0,n)
# state = c(S,I_1,I_2,...I_n)   
    S = state[1]
    I[1:n] = state[2:(n+1)]
    X = sum(I[(m+1):n])
    dSdt = mu - beta*X*S - mu*S
    dIdt[1] = beta*X*S - gamma*n*I[1] - mu*I[1]
    for (i in 2:n){
        dIdt[i] = gamma*n*I[i-1] - gamma*n*I[i] - mu*I[i]
    }
    return(list(c(dSdt,dIdt)))
  })
}

# Set parameters (all rates are per day)
n=13           # Number of stages in the infected period
m=8            # Number of stages in the exposed period
gamma=1/13     # Removal or recovery rate
beta=17/5      # Transmission rate
mu=1/(55*365)  # Death rate (assumed nu=mu)
pars = list(n,m,beta,gamma,mu)

# Initial conditions
S0=0.05       # Initial proportion of population that is susceptible
I0=0.00001    # Initial proportion of population infected (I_i(0) = I0/n)

# Time range
MaxTime=30*365
step=1
times = seq(0,MaxTime,step)

# Check all parameters are valid
if(I0<=0){
  print(paste("Initial level of infection is less than or equal to zero, I0 =",I0))  
  stop()}
if(S0<=0){
  print(paste("Initial level of susceptibles is less than or equal to zero, S0 =",S0))  
  stop()}
if(n<=0){
  print(paste("Number of infected stages is less than or equal to zero, n =",n))  
  stop()}
if(m>=n){
  print(paste("Number of infectious stages is less than or equal to zero, m-n =",m-n))  
  stop()}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(beta<=0){
  print(paste("Transmission rate is less than or equal to zero, beta =",beta))  
  stop()}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}  

I0=I0*rep(1,n)/n     # Vector of initial proportion infected
y0=c(S0,I0)

# Solve differential equations
library(deSolve)
SEIR_mult.out = ode(y0,times,SEIR_mult,pars,method="ode45",rtol=1e-5)

out = as.data.frame(SEIR_mult.out)
# create column names
p = c(1:n); q = as.character(p)
labels = paste("I",q,sep="")
colnames(out) = c("time","S",labels)

library(ggplot2)
p1=ggplot() +
  geom_line(data=out,aes(y=S,x= time),colour="green2")+
  xlab("Time") + ylab("Susceptibles, S") +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black",
                                    fill = NA, linewidth = 0.6)) 

# create long vector for this plot
library(tidyr)
out_long=pivot_longer(out,cols=c(all_of(labels)),
                       names_to="inf_group",values_to="infectious")

E_colour = rep("orange",m)
I_colour = rep("red",(n-m))
p2=ggplot() +
  geom_line(data=out_long,aes(y=infectious,x= time,group=inf_group,colour=inf_group))+
  xlab("Time") + ylab(bquote("Infection, I"[i])) +
  scale_color_manual(name = "",values=c(E_colour,I_colour)) +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = "none") +
  theme(panel.border = element_rect(color = "black",
                                    fill = NA, linewidth = 0.6))

if(n>1){
p3=ggplot() +
  geom_line(data=out,aes(y=rowSums(out[,c(3:(n+2))]),x= time),colour="red")+
  xlab("Time") + ylab("Total infection") +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black",
                                    fill = NA, linewidth = 0.6)) 
}else{
  p3=ggplot() +
    geom_line(data=out,aes(y=out[,3],x= time),colour="red")+
    xlab("Time") + ylab("Total infection") +
    scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
    theme_classic() +
    theme(panel.border = element_rect(color = "black",
                                      fill = NA, linewidth = 0.6))  
}

library(patchwork)
p1/p2/p3