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

# It is a 2 strain SIR disease model, where parameter alpha governs partial
# susceptibility and parameter a governs partial transmissibility. Note
# that in theory alpha or a could be greater than one to imply enhanced
# susceptibility or transmissibility.

# Set up differential equations as a function
SIR_2strain = function(time,state,pars){
  with(as.list(c(state,pars)),{
    
    lambda1 = beta[1]*(NIS+a[1]*NIR)
    lambda2 = beta[2]*(NSI+a[2]*NRI)
    
    dNSSdt = mu - NSS*(lambda1 + lambda2)  - mu*NSS
    dNISdt = NSS*lambda1 - gamma[1]*NIS - mu*NIS
    dNRSdt = gamma[1]*NIS - alpha[2]*NRS*lambda2 - mu*NRS
    
    dNRIdt = alpha[2]*NRS*lambda2 - gamma[2]*NRI - mu*NRI
    dNSIdt = NSS*lambda2 - gamma[2]*NSI - mu*NSI
    
    dNSRdt = gamma[1]*NSI - alpha[1]*NSR*lambda1 - mu*NSR
    dNIRdt = alpha[1]*NSR*lambda1 - gamma[1]*NIR - mu*NIR
    
    return(list(c(dNSSdt,dNISdt,dNRSdt,dNRIdt,dNSIdt,dNSRdt,dNIRdt)))
  })
}

# Checks whether parameters are all positive
CheckPositive = function(p){
  n = any(p<0)
  if(n==TRUE){
    print(paste("Negative value in input parameter"))
    stop()}
}

# Set parameters (all rates are per day)
beta=c(260, 520)/365.0    # beta[i] is transmission rate for strain i
gamma=c(1.0, 1.0)/7.0     # gamma[i] is recovery rate for strain i  
mu = 1/(70*365.0)         # per capita death rate (equal to birth rate)
alpha = c(0.5, 0.4)       # alpha[i] is the modified susceptibility to strain i for individuals recovered from the other strain
a = c(0.4, 0.5)           # a[i] is the modified transmission rate of strain i from those that have recovered from the other strain
pars = c(beta,gamma,mu,alpha,a)

# Check sizes for terms that are vectors
library(vctrs)
vec_check_size(beta,size=2)
vec_check_size(gamma,size=2)
vec_check_size(alpha,size=2)
vec_check_size(a,size=2)

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

# Checks all parameters are valid
CheckPositive(beta)
CheckPositive(gamma)
CheckPositive(alpha)
CheckPositive(a)
CheckPositive(mu)
CheckPositive(MaxTime)

# Initial proportion of population in each group - set to suit default parameters
NSS0 = 0.1     
NIS0= 1e-04       
NRS0 = 0.02
NRI0 = 0.0
NSI0 = 1e-04
NSR0 = 0.5
NIR0 = 0.0
y0 = c(NSS=NSS0,NIS=NIS0,NRS=NRS0,NRI=NRI0,NSI=NSI0,NSR=NSR0,NIR=NIR0)

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

out = as.data.frame(SIR_2strain.out)
out$NRR = 1.0 - rowSums(out[,2:8])
# I_1 = NIS + a1*NIR
out$inf1 = out[,3]+a[1]*out[,8]
# I_2 = NSI + a2*NRI
out$inf2 = out[,6]+a[2]*out[,5]

# Create long data frames for plotting
library(tidyr)
out_long=pivot_longer(out,cols=c("NSS","NRS","NSR","NRR"),
                      names_to="strain_group",values_to="uninfected")
out_long1=pivot_longer(out,cols=c("NIS","NIR","inf1"),
                      names_to="strain_group",values_to="infectious1")
out_long2=pivot_longer(out,cols=c("NSI","NRI","inf2"),
                       names_to="strain_group",values_to="infectious2")

# Plot results using ggplot2
library(ggplot2)
p1=ggplot(data=out_long,aes(x=time/1e4,y=uninfected,group=strain_group,colour=strain_group)) +
  geom_line() +
  xlab(bquote("Time " (x10^4))) + ylab("Uninfected") +
  scale_color_manual(name = "", breaks=c("NSS","NRS","NSR","NRR"),
                     labels = c(expression(N["SS"]),expression(N["RS"]),expression(N["SR"]),expression(N["RR"])),
                     values = c("blue2","red","orange","purple3")) +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.93,0.7),
        legend.key.size=unit(3,"mm"),
        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(data=out_long1,aes(x=time/1e4,y=infectious1/1e-04,group=strain_group,colour=strain_group)) +
  geom_line() +
  xlab(bquote("Time " (x10^4))) + ylab(bquote("Infectious 1 "(x10^-4))) +
  scale_color_manual(name = "", breaks=c("NIS","NIR","inf1"),
                     labels = c(expression(N["IS"]),expression(N["IR"]),expression(I["1"])),
                     values = c("blue2","red","orange")) +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  scale_y_continuous(limits=c(0,3)) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.93,0.75),
        legend.key.size=unit(3,"mm"),
        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))
  
p3=ggplot(data=out_long2,aes(x=time/1e4,y=infectious2/1e-04,group=strain_group,colour=strain_group)) +
  geom_line() +
  xlab(bquote("Time " (x10^4))) + ylab(bquote("Infectious 2 " (x10^-4))) +
  scale_color_manual(name = "", breaks=c("NSI","NRI","inf2"),
                     labels = c(expression(N["SI"]),expression(N["RI"]),expression(I["2"])),
                              values = c("blue2","red","orange")) +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.93,0.75),
        legend.key.size=unit(3,"mm"),
        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