# This is the R version of program 4.4 from page 136 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is a simple model of vector borne transmission of infection. We
# implicitly assume a vector-notation, with humans being the first entry
# and mosquitoes being the second.
# In the default parameterisation of the model mosquitoes never recover 
# from infection, although this is allowed by the general formulation.

# Set up differential equations as a function
SIR_vectors = function(time,state,pars){
  with(as.list(c(state,pars)),{
    X = state[1:2]
    Y = state[3:4]
    dXdt= rep(0,2)
    dYdt = rep(0,2)
    
    dXdt[1] = nu[1] - r*(T[1,2]*Y[2] + T[1,1]*Y[1])*X[1] - mu[1]*X[1]  # X_H
    dXdt[2] = nu[2] - r*(T[2,2]*Y[2] + T[2,1]*Y[1])*X[2] - mu[2]*X[2]  # X_M
    
    dYdt[1] = r*(T[1,2]*Y[2] + T[1,1]*Y[1])*X[1] - gamma[1]*Y[1] - mu[1]*Y[1]  # Y_H
    dYdt[2] = r*(T[2,2]*Y[2] + T[2,1]*Y[1])*X[2] - gamma[2]*Y[2] - mu[2]*Y[2]  # Y_M
    
    return(list(c(dXdt,dYdt)))
  })
}

# 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)
r=0.5/1e3   # rate at which humans are bitten
T=matrix(c(0, 0.5, 0.8, 0), 
            nrow = 2, ncol = 2, byrow = TRUE) 
              # T_ij is the transmission probability (following a bite) to species i from species j
gamma=c(0.033, 0)   # gamma_i is the recovery rate for host species i
mu=c(5.5e-5, 0.143) # mu_i is per capita death rate for host species i
nu=c(5.5e-2, 1.443e3) # nu_i is the birth rate for host species i
X0 = c(1e3, 1e4)      # X0_i is initial number of individuals of species i that are susceptible
Y0 = c(1,1)           # Y0_i is the number of individuals of sepcies i that are infectious
pars = c(r,T,gamma,mu,nu)

# Check sizes for terms that are vectors/matrices
library(vctrs)
vec_check_size(mu,size=2)
vec_check_size(nu,size=2)
vec_check_size(X0,size=2)
vec_check_size(Y0,size=2)
mat_size = dim(T)
if(mat_size[1]!=2 | mat_size[2]!=2) {
  print(paste("Transmission matrix not of correct size"))  
  stop()}  

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

# Checks all parameters are valid
CheckPositive(r)
CheckPositive(gamma)
CheckPositive(mu)
CheckPositive(nu)
CheckPositive(X0)
CheckPositive(Y0)
CheckPositive(MaxTime)
if(T[1,1]!=0 | T[2,2]!=0) {
  print(paste("Transmission probability within species (the diagonal of T) should be zero"))  
  stop()}
if(T[1,2]<0 | T[1,2]>1) {
  print(paste("Transmission probability to humans from mosquitos is not between 0 and 1"))  
  stop()} 
if(T[2,1]<0 | T[2,1]>1) {
  print(paste("Transmission probability to mosquitos from humans is not between 0 and 1"))  
  stop()} 

# Initial conditions set to suit default parameters
v0 = c(X=X0,Y=Y0)

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

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

# Plot the results using ggplot2
library(ggplot2)

p1=ggplot() +
  geom_line(data=out,aes(y=X1,x= time),linetype=1,color="green3") +
  xlab("Time") + ylab("Susceptible humans") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  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=X2,x= time),linetype=1,color="green3") +
  xlab("Time") + ylab("Susceptible mosquitos") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  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=Y1,x= time),linetype=1,color="red") +
  xlab("Time") + ylab("Infected humans") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))

p4=ggplot() +
  geom_line(data=out,aes(y=Y2,x= time),linetype=1,color="red") +
  xlab("Time") + ylab("Infected mosquitos") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))

library(patchwork)
(p1|p2)/(p3|p4)