# This is the R version of program 4.2 from page 123 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is a N strain SIR disease model, where the strains are arranged in a
# circle and each strain offers partial immunity (in terms of reduced
# transmission) to its neighbours.

# You may also wish to try parameters:
#  (N,beta,gamma,mu,a,S0,P0,lambda0,MaxTime) =  
#  (4,40,9.98,0.02,0.25,[0.25 0.14 0.25 0.14],[0.016 0.55 0.016 0.55],
#   [0.07 1e-12 0.07 1e-12],200)

# Set up differential equations as a function
SIR_Nstrain = function(time,state,pars){
  with(as.list(c(state,pars)),{
    S = state[1:N]
    P = state[(N+1):(2*N)]
    lambda = state[(2*N+1):(3*N)]
    dSdt= rep(0,N)
    dPdt = rep(0,N)
    dlambdadt = rep(0,N)
    for (i in 1:4){
      r = 1 + i%%N
      l = 1 + (i+N-2)%%N
    dSdt[i] = mu - S[i]*(lambda[i]+lambda[l]+lambda[r]) - mu*S[i]
    dPdt[i] = S[i]*(lambda[l]+lambda[r]) - P[i]*lambda[i] - mu*P[i]
    dlambdadt[i] = beta*(S[i]+a*P[i])*lambda[i] - gamma*lambda[i] - mu*lambda[i]
    }   
    return(list(c(dSdt,dPdt,dlambdadt)))
  })
}

# 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)
N=4        # Number of strains
beta=40.0  # Transmission rate (same for all strains)
gamma=9.98 # Recovery rate (same for all strains)
mu=0.02    # Per capita death rate
a=0.4      # Modified transmission rate due to partial immunity
S0=c(0.08, 0.1, 0.1, 0.11)  # S_i is the initial proportion of population that are susceptible to strain i
P0=c(0.4, 0.3, 0.3, 0.29)   # P_i is the initial proprtion of population that are partially immune to strain i
lambda0=c(0.15, 0.02, 0.1, 0.01)  # lambda_i is the initial force of infection due to strain i
pars = c(N,beta,gamma,mu,a)

# Check sizes for terms that are vectors
library(vctrs)
vec_check_size(S0,size=N)
vec_check_size(P0,size=N)
vec_check_size(lambda0,size=N)

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

# Checks all parameters are valid
CheckPositive(beta)
CheckPositive(gamma)
CheckPositive(a)
CheckPositive(mu)
CheckPositive(S0)
CheckPositive(P0)
CheckPositive(lambda0)
CheckPositive(MaxTime)

# Initial conditions set to suit default parameters
y0 = c(S=S0,P=P0,lambda=lambda0)

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

out = as.data.frame(SIR_Nstrain.out)

# Create long data frames for plotting
 library(tidyr)
 out_long=pivot_longer(out,cols=c("S1","S2","S3","S4"),
                       names_to="strain_group",values_to="susceptibles")
 out_long1=pivot_longer(out,cols=c("P1","P2","P3","P4"),
                       names_to="strain_group",values_to="partial")
 out_long2=pivot_longer(out,cols=c("lambda1","lambda2","lambda3","lambda4"),
                       names_to="strain_group",values_to="force")

# Plot results using ggplot2
library(ggplot2)
 
p1=ggplot(data=out_long,aes(x=time,y=susceptibles,group=strain_group,colour=strain_group)) +
   geom_line() +
   xlab("Time") + ylab("Susceptibles") +
   scale_color_manual(name = "", breaks=c("S1","S2","S3","S4"),
                      labels = c(expression(S["1"]),expression(S["2"]),expression(S["3"]),expression(S["4"])),
                      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.94,0.3),
         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,y=partial,group=strain_group,colour=strain_group)) +
  geom_line() +
  xlab("Time") + ylab("Partially susceptible") +
  scale_color_manual(name = "", breaks=c("P1","P2","P3","P4"),
                     labels = c(expression(P["1"]),expression(P["2"]),expression(P["3"]),expression(P["4"])),
                     values = c("blue2","red","orange","purple3")) +
  scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
  scale_y_continuous(limits=c(0.1,0.4)) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.94,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))


p3=ggplot(data=out_long2,aes(x=time,y=force,group=strain_group,colour=strain_group)) +
  geom_line() +
  xlab("Time") + ylab("Force of infection") +
  scale_color_manual(name = "", breaks=c("lambda1","lambda2","lambda3","lambda4"),
                     labels = c(expression(lambda["1"]),expression(lambda["2"]),expression(lambda["3"]),expression(lambda["4"])),
                     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.94,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))


 library(patchwork)
 p1/p2/p3