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

# Set up equations as a function
SIS_risk = function(time,state,pars){
  with(as.list(c(state,pars)),{
    I=c(I1,I2,I3,I4,I5)
    S=n-I
    dIdt = (beta%*%I)*S - gamma*I
    return(list(c(dIdt)))
    })
}

# Set parameters (all rates are per year)
m=5                      # number of risk groups
b=c(0, 3, 10, 60, 100)
beta=0.0016*b%*%t(b)     # Matrix of transmission rates
gamma=0.2*c(1,1,1,1,1)   # Vector of recovery rate
n=c(0.06,0.31,0.52,0.08,0.03)   # Vector of proportion of population that are in each risk group
parms = list(beta,gamma,n)

# Initial conditions
I0 = c(I1=0.0, I2=0.0, I3=0.0, I4=0.0, I5=1e-5)   # Vector of initial proportions of population 
                                                  # that are both infectious and in each risk group
# Time range
MaxTime=30
step=0.1
times = seq(0,MaxTime,step)

# Check size of all inputs
A=dim(beta)
if(A[1]!=m | A[2]!=m){
  print(paste("Size of beta matrix does not correspond to number of risk groups"))  
  stop()}
if(length(gamma)!=m){
  print(paste("Size of gamma vector does not correspond to number of risk groups"))  
  stop()}
if(length(n)!=m){
  print(paste("Size of n vector does not correspond to number of risk groups"))  
  stop()}
if(length(I0)!=m){
  print(paste("Size of I vector does not correspond to number of risk groups"))  
  stop()}

# Check all parameters are valid
for (i in 1:m) {
  if(I0[i]<0){
  print(paste("Initial level of infection in risk group is less than zero, I0[",i,"] =",I0[i]))   
  stop()}
  if(n[i]<=0){
    print(paste("A risk group is less than or equal to zero, n[",i,"] =",n[i]))   
    stop()}  
  if(I0[i]>n[i]){
    print(paste("Initial level of infection in risk group is greater than the size of the group, n[",i,"] =",n[i],", I0[",i,"]=",I0[i]))   
    stop()}   
  if(gamma[i]<=0){
    print(paste("Recovery rate is less than or equal to zero, gamma[",i,"] =", gamma[i]))   
    stop()} 
  for (j in 1:m){ 
    if(beta[i,j]<0){
      print(paste("Transmission rate beta is less than zero, i,j,beta[i,j]=",i,j,beta[i,j]))
      stop()}
    }
}

if(abs(sum(n)-1)>0.1){
   print(paste("Sum of the risk groups does not equal one, sum(n)=",sum(n)))
    stop()}
if(abs(sum(n)-1)>0.01){
  print(paste("Warning:Sum of the risk groups does not equal one: n and I have been rescaled, sum(n)=",sum(n)))
  n=n/sum(n)
  I0=I0/sum(n)}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime=",MaxTime))
  stop()}
if(length(which(abs(beta-t(beta))>0.001))>0){
    print(paste("Transmission rate beta is not symmetric"))
    stop()}

# Solve differential equations
library(deSolve)
SIS_risk.out = ode(I0,times,SIS_risk,parms,method="ode45",rtol=1e-5)

out = as.data.frame(SIS_risk.out)
colnames(out) = c("time","I1","I2","I3","I4","I5")

# Create long data frame for plotting
library(tidyr)
out_long=pivot_longer(out,cols=c("I1","I2","I3","I4","I5"),
             names_to="risk_group",values_to="infectious")

# Plot results using ggplot2
library(ggplot2)
p1=ggplot(data=out_long,aes(x=time,y=infectious,group=risk_group,colour=risk_group)) +
     geom_line() +
     xlab("Time") + ylab("Infectious, I") +
     scale_color_manual(name = "", 
          labels=c("Risk group 1","Risk group 2","Risk group 3","Risk group 4","Risk group 5"),
          values = c("green","green3","green4","red4","red")) +
     scale_x_continuous(expand=expansion(mult=c(0,0.0)),limits=c(0,30)) +
     theme_classic() +
     theme(legend.title = element_blank(),legend.position = c(0.25,0.8),
          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_long,aes(x=time,y=infectious,group=risk_group,colour=risk_group)) +
    geom_line()+
    xlab("Time") + ylab("Infectious, I") +
    scale_color_manual(name = "", 
           labels=c("Risk group 1","Risk group 2","Risk group 3","Risk group 4","Risk group 5"),
           values = c("green","green3","green4","red4","red")) +
    scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
    scale_x_continuous(expand=expansion(mult=c(0,0)),limits=c(0,30)) +
    theme_classic() +
    theme(legend.title = element_blank(),legend.position = c(0.75,0.2),
          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_long,aes(x=time,y=infectious/n,group=risk_group,colour=risk_group)) +
  geom_line() +
  xlab("Time") + ylab("Infectious, I") +
  scale_color_manual(name = "", 
                     labels=c("Risk group 1","Risk group 2","Risk group 3","Risk group 4","Risk group 5"),
                     values = c("green","green3","green4","red4","red")) +
  scale_x_continuous(expand=expansion(mult=c(0,0.0)),limits=c(0,30)) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.25,0.8),
        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))

p4=ggplot(data=out_long,aes(x=time,y=infectious/n,group=risk_group,colour=risk_group)) +
  geom_line()+
  xlab("Time") + ylab("Infectious, I") +
  scale_color_manual(name = "", 
                     labels=c("Risk group 1","Risk group 2","Risk group 3","Risk group 4","Risk group 5"),
                     values = c("green","green3","green4","red4","red")) +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  scale_x_continuous(expand=expansion(mult=c(0,0)),limits=c(0,30)) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.75,0.2),
        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 | p4)