# This is the R version of program 5.4 from page 186 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is a model for Rabit Hemorrhagic Disease, in which both transmission
# rate and birth rates can be seasonally forced.
# Note that we are using numbers of individuals in this formulation.
# Bifurcation plots are not possible with this code.

# Set up differential equations as a function
Rabbit_HD = function(t,state,pars){
  with(as.list(c(state,pars)),{
    
    alpha = alpha0*(1+alpha1*sin(2.0*pi*t/365.0))
    beta = beta0*(1+beta1*sin(2.0*pi*t/365.0))
    
    dXdt = alpha*N - beta*X*Y - (mu+N/K)*X
    dYdt = beta*X*Y - (gamma + mu + m + N/K)*Y
    dNdt = (alpha - mu - N/K)*N - m*Y
    
    return(list(c(dXdt,dYdt,dNdt)))
  })
}

# Set parameters (all rates are per day)
beta0=0.936   # Mean transmission rate
beta1=0.1     # Amplitude of sinusoidal forcing
gamma=0.025   # Recovery rate
alpha0=0.02   # Mean birth rate
alpha1=0.1    # Amplitude of sinusoidal forcing for birth rate
mu=0.01       # Per capita death rate
m=0.475       # Mortality rate due to infection
K=10000       # Carrying capacity accociated with the host populations
pars = c(beta0,beta1,gamma,alpha0,alpha1,mu,m,K)

# Initial conditions
X0=0.5        # Initial number of susceptible hosts (rabbits)
Y0=0.01       # Initial number of infectious hosts (rabbits)
N0=0.6        # Initial total population size
y0 = c(X=X0,Y=Y0,N=N0)

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

# Check all parameters are valid
if(X0<=0){
  print(paste("Initial number of susceptibles is less than or equal to zero, X0 =",X0))  
  stop()}
if(Y0<=0){
  print(paste("Initial number of infecteds is less than or equal to zero, Y0 =",Y0))  
  stop()}
if(N0<=0){
  print(paste("Initial population size is less than or equal to zero, N0 =",N0))  
  stop()}
if(beta0<=0){
  print(paste("Transmission rate beta0 is less than or equal to zero, beta0 =",beta0))  
  stop()}
if(beta1<0){
  print(paste("Seasonality beta1 is less than zero, beta1 =",beta1))  
  stop()}
if(beta1>=1){
  print(paste("Seasonality beta1 is geater than or equal to 1, beta1 =",beta1))  
  stop()}
if(alpha1<0){
  print(paste("Seasonality alpha1 is less than zero, alpha1 =",alpha1))  
  stop()}
if(alpha1>=1){
  print(paste("Seasonality alpha1 is geater than or equal to 1, alpha1 =",alpha1))  
  stop()}
if(gamma<=0){
  print(paste("Recovery rate gamma is less than or equal to zero, gamma =",gamma))  
  stop()}
if(alpha0<0){
  print(paste("Birth rate alpha0 is less than zero, alpha0 =",alpha0))  
  stop()}
if(mu<0){
  print(paste("Death rate mu is less than zero, mu =",mu))  
  stop()}
if(m<0){
  print(paste("Mortality rate due to infection m is less than zero, m =",m))  
  stop()}
if(K<=0){
  print(paste("Carrying capacity K is less than or equal to zero, K =",K))  
  stop()}
if(MaxTime<=0){
  print(paste("Max run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if(X0+Y0>N0){
  print(paste("Initial level of susceptibles+infecteds is greater than the population size, X0+Y0 =",X0+Y0,", N0 = ",N0))  
  stop()}
if(beta0<(gamma+mu)){
  print(paste("Basic reproductive ratio is less than 1, beta0/(gamma+mu) =",beta0/(gamma+mu)))  
  stop()}


# The main iteration
# Solve differential equations defined in function Rabbit_HD
library(deSolve)
Rabbit_HD.out = ode(y0,times,Rabbit_HD,pars,rtol=1e-5)
out = as.data.frame(Rabbit_HD.out)

# Plot the results using ggplot2
library(ggplot2)

p1=ggplot() +
  geom_line(data=out,aes(y=X,x= time/365),
            linewidth=0.3,linetype=1,color="green3") +
  xlab("Time (years)") + ylab("Susceptibles") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),limits=c(0.4,0.7)) +
  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=Y,x= time/365),
            linewidth=0.3,linetype=1,color="red") +
  xlab("Time (years)") + ylab("Infectious") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),limits=c(0.005,0.03)) +
  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=N,x= time/365),
            linewidth=0.3,linetype=1,color="black") +
  xlab("Time (years)") + ylab("Population size") +
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),limits=c(0.5,0.7)) +
  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