# This is the R version of program 2.3 from page 35 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.

# It is the SIR model with a probability of mortality, and
# unequal births and deaths. This code assumes Density-
# Dependent Transmission

# Set up equations as a function
XYZ_dd = function(time,state,pars){
  with(as.list(c(state,pars)),{
    dXdt =  nu - beta*X*Y - mu*X
    dYdt = beta*X*Y - (gamma+mu)*Y/(1-rho)
    dZdt = gamma*Y - mu*Z
    return(list(c(dXdt,dYdt,dZdt)))
  })
}

# Set parameters (all rates are per day)
beta=520/365.0    # Transmission rate 
gamma=1/7.0       # Recovery rate
mu=1/(70*365.0)   # Per capita death rate
nu=1/(70*365.0)   # Population level birth rate
rho=0.5           # Mortality probability
pars = c(beta,gamma,mu,nu,rho)

# Initial conditions
X0 = 0.2       # initial number or density of susceptible individuals
Y0 = 1e-4      # initial number or density of infectious individuals
N0 = 1           # total population size
Z0 = N0-X0-Y0
v0 = c(X=X0,Y=Y0,Z=Z0)

# Time range
MaxTime=1e5
step=28
times = seq(0,MaxTime,step)

# Check all parameters are valid
if(X0<=0){
  print(paste("Initial level of susceptibles is less than or equal to zero, X0 =",X0))  
}
if(Y0<=0){
  print(paste("Initial level of infecteds is less than or equal to zero, Y0 =",Y0))  
}
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(mu<0){
  print(paste("Birth/death rate is less than zero, mu =",mu))  
}
if(nu<0){
  print(paste("Birth/death rate is less than zero, nu =",nu))  
}
if(rho<0 | rho>1){
  print(paste("Mortality probability is not a valid probability, rho =",rho))  
}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
}
if(X0+Y0>N0){
  print(paste("Initial level of susceptibles+infecteds is greater than population size, X0+Y0  =",X0+Y0, "N0 =",N0))  
}
if(beta*(1-rho)*nu<(gamma+mu)*mu){
  print(paste("Basic reproductive ratio is less than 1, R =",beta*(1-rho)*nu/((gamma+mu)*mu)))
}

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

# Plot the results using ggplot2
library(ggplot2)

p1=ggplot() +
  geom_line(data=out,aes(y=X,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=out,aes(y=Y,x= time/1e4),linetype=1,color="red") +
  xlab(bquote("Time " (x10^4))) + 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=Z,x= time/1e4,linetype="Y1"),
            color="black") +
  geom_line(data=out,aes(y=X+Y+Z,x= time/1e4,linetype="Y2"),
            color="black") +
  xlab(bquote("Time " (x10^4))) + ylab("Recovereds/Population") +
  scale_linetype_manual(name = "", labels=c("Recovereds","Total population"),
                        values = c(1,2)) +  
  scale_y_continuous(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.85,0.7),
            legend.box.background = element_rect(colour = "black"),
            legend.spacing.y = unit(0, "mm")) +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))

library(patchwork)
p1/p2/p3