# This is the R version of program 3.1 from page 58 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SIS model with two different risk-groups.
# Note we no longer explicitly model the susceptible class.

# Set up equations as a function
SIS_2 = function(time,state,pars){
  with(as.list(c(state,pars)),{
    SH=nH-IH
    SL=(1-nH)-IL
    dIHdt = (beta[1,1]*IH+beta[1,2]*IL)*SH - gamma*IH
    dILdt = (beta[2,1]*IH+beta[2,2]*IL)*SL - gamma*IL
    return(list(c(dIHdt,dILdt)))
  })
}

# Set parameters (all rates are per year)
beta=matrix(c(10, 0.1, 0.1, 1), 
            nrow = 2, ncol = 2, byrow = TRUE)    # Matrix of transmission rates
gamma=1.0       # Recovery rate
nH=0.2          # Proportion of population that are high risk
pars = c(beta,gamma,nH)

# Initial conditions
IH0 = 1e-5      # Initial proportion of population that are both infectious and high risk
IL0 = 1e-3      # Initial proportion of population that are both infectious and low risk
y0 = c(IH=IH0,IL=IL0)

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

# Check all parameters are valid
if(IH0<0){
  print(paste("Initial level of infection in high-risk group is less than zero, IH0 =",IH0))  
  stop()}
if(IH0>nH){
  print(paste("Initial level of infection in high-risk group is greater than the high-risk population size, IH0 =",IH0, ", nH =",nH))  
  stop()}
if(IL0<0){
  print(paste("Initial level of infection in low-risk group is less than zero, IL0 =",IL0))   
  stop()}
if(IL0>(1-nH)){
  print(paste("Initial level of infection in low-risk group is greater than the low-risk population size, IL0 =", IL0, ", nL =",1-nH))    
  stop()}
if(IH0+IL0<=0){
  print(paste("Initial level of total infection is less than or equal to zero, IH0+IL0 =",IH0+IL0))  
  stop()}
if(nH<=0){
  print(paste("High-risk population is less than or equal to zero, nH =",nH))  
  stop()}
if(nH>=1){
  print(paste("Low-risk population is less than or equal to zero, nL =",1-nH))  
  stop()}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if(beta[1,2]!=beta[2,1]){
  print(paste("Transmission rate beta is not symmetric beta =",beta[1,1],beta[1,2],beta[2,1],beta[2,2]))  
  stop()}
if(length(which(beta<0))>0){
print(paste("Transmission rate is less than zero, beta =",beta[1,1],beta[1,2],beta[2,1],beta[2,2]))
stop()}

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

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

p1=ggplot() +
  geom_line(data=out,aes(y=IH,x= time,linetype="Y1"),
            color="red") +
  geom_line(data=out,aes(y=IL,x= time,linetype="Y2"),
            color="red") +
  xlab("Time") + ylab("Infectious") +
  scale_linetype_manual(name = "", labels=c("High risk","Low risk"),
                        values = c(1,2)) +  
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),limits=c(0,0.12)) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.13,0.83),
        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))

p2=ggplot() +
  geom_line(data=out,aes(y=IH,x= time,linetype="Y1"),
            color="red") +
  geom_line(data=out,aes(y=IL,x= time,linetype="Y2"),
            color="red") +
  xlab("Time") + ylab("Infectious") +
  scale_linetype_manual(name = "", labels=c("High risk","Low risk"),
                        values = c(1,2)) +  
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.87,0.18),
        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