# This is the R version of program 7.2 from page 242 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SIR epidemic in a metapopulationFor simplicity births and deaths 
# have been ignored, and we work with numbers of individuals.
# Y[i][j] refers to infected individual who are currently in i but live in j.

# Set up equations as a function
SIR_meta2 = function(time,state,pars){
  with(as.list(c(state,pars)),{
    nsq=n*n
    X = matrix(state[1:nsq],nrow = n, ncol = n, byrow = FALSE)
    dXdt = matrix(c(rep(0,nsq)),nrow=n,ncol=n)
    
    ystart=nsq+1; yend = 2*nsq
    Y = matrix(state[ystart:yend],nrow = n, ncol = n, byrow = FALSE)
    dYdt = matrix(c(rep(0,nsq)),nrow=n,ncol=n)
    
    Nstart=2*nsq+1; Nend = 3*nsq
    N = matrix(state[Nstart:Nend],nrow = n, ncol = n, byrow = FALSE)
    dNdt = matrix(c(rep(0,nsq)),nrow=n,ncol=n)
    
    sumY=rep(0,n); sumN=rep(0,n)
    for (i in 1:n){
      for (j in 1:n){
        sumY[i] = sumY[i] + Y[i,j]
        sumN[i] = sumN[i] + N[i,j]
      }
    }
 
    for (i in 1:n){
      for (j in 1:n){
        if(i!=j){ 
          dXdt[i,j] = -beta[i]*X[i,j]*sumY[i]/sumN[i] + l[i,j]*X[j,j] - r[i,j]*X[i,j]
          dYdt[i,j] = beta[i]*X[i,j]*sumY[i]/sumN[i] - gamma[i]*Y[i,j] + l[i,j]*Y[j,j] - r[i,j]*Y[i,j]
          dNdt[i,j] = l[i,j]*N[j,j] - r[i,j]*N[i,j]
        } 
      }
          dXdt[i,i] = -beta[i]*X[i,i]*sumY[i]/sumN[i] - sum(l[,i])*X[i,i]
          dYdt[i,i] = beta[i]*X[i,i]*sumY[i]/sumN[i] - gamma[i]*Y[i,i] - sum(l[,i])*Y[i,i]
          dNdt[i,i] = -sum(l[,i])*N[i,i]
    }
      
    for (i in 1:n){
      for (j in 1:n){
        dXdt[i,i] = dXdt[i,i] + r[j,i]*X[j,i]
        dYdt[i,i] = dYdt[i,i] + r[j,i]*Y[j,i]
        dNdt[i,i] = dNdt[i,i] + r[j,i]*N[j,i]
      }
    }
    
    return(list(c(dXdt,dYdt,dNdt)))
  })
}

# Set parameters (all rates are per day)
n = 5               # Number of subpopulations
beta = rep(1.0,n)    # Transmission rate for each subpopulation
gamma = rep(0.3,n)   # Recovery rate for each subpopulation
n2=n*n
l=matrix(c(rep(0,n2)),nrow=n, ncol=n, byrow=TRUE)
for (i in 1:n){
  for (j in 1:n){
    if(abs(i-j)==1){
      l[i,j]=0.1
    }
  }
}
r = matrix(c(rep(2.0,n2)),nrow = n, ncol = n, byrow = TRUE)
r = r - diag(2.0,n,n)

pars = list(beta,gamma,l,r,n)

# Initial conditions
N0=rep(1000,n)   # Initial number of individuals in each spatial class
X0=rep(800,n)   # Initial number of susceptible individuals in each spatial class
Y0=rep(0.0,n)   # Initial number of infectious individuals in each spatial class
Y0[1] = 1
X=diag(X0); Y=diag(Y0); NN=diag(N0)
v0 = c(X,Y,NN)

MaxTime=60
step=0.1
T0=0

# Check all parameters are valid
  for (i in 1:n) {
    if(X0[i]<0){
      print(paste("Initial level of susceptibles is less than zero, X0[",i,"] =",X0[i]))   
      stop()}
    if(Y0[i]<0){
     print(paste("Initial level of infectious is less than zero, Y0[",i,"] =",Y0[i]))   
      stop()}  
    if(N0[i]<0){
      print(paste("Initial number of individuals is less than zero, N0[",i,"] =",N0[i]))   
      stop()} 
    if(X0[i]+Y0[i]>N0[i]){
      print(paste("Initial level of susceptibles+infecteds is greater than population size, X0+Y0[",i,"] =",X0[i]+Y0[i])) 
      stop()} 
    for (j in 1:n) {
      if(l[i,j]<0){
        print(paste("Movement rate is less than zero, l[",i,",",j,"] =",l[i,j])) 
        stop()} 
      if(r[i,j]<0) {
      print(paste("Return rate is less than zero, r[",i,",",i,"] =",r[i,j])) 
      stop()}  
    }
  }
  if(length(which(beta<0))>0){
    print(paste("Transmission rate beta has a term that is less than zero"))
    stop()} 
 if(length(which(gamma<0))>0){
    print(paste("Recovery rate has a term that is less than zero,"))  
    stop()}
  if(MaxTime<=0){
    print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
    stop()}

# Solve differential equations
times =seq(T0,MaxTime,step)
  
  library(deSolve)
  SIR_meta2.out = ode(v0,times,SIR_meta2,pars,method="ode45",rtol=1e-4)

  outX=matrix(c(rep(0.0,(n*(1+MaxTime/step)))),nrow=(1+MaxTime/step),ncol=n)
  outY=matrix(c(rep(0.0,(n*(1+MaxTime/step)))),nrow=(1+MaxTime/step),ncol=n)

  # For each subpopulation i, sum over subpopulation j
  for (t in 1:(1+MaxTime/step)){
    for (i in 1:n){
      for (k in 1:n){
      j= (i+1) + (k-1)*n
      l= (n*n+1) + i + (k-1)*n
      outX[t,i] = outX[t,i]+SIR_meta2.out[t,j]
      outY[t,i] = outY[t,i]+SIR_meta2.out[t,l]
      }
    }
  }
  
  out = as.data.frame(cbind(SIR_meta2.out[,1],outX,outY)) 
  colnames(out) = c("time",paste0("X", 1:n),paste0("Y",1:n))
  
  # Create long data frames for plotting
  library(tidyr)
  n1=n+1; n2=n+2; nend=2*n+1
  out_long1=pivot_longer(out,cols=all_of(c(2:n1)),
                         names_to="subpop",values_to="susceptible")
  out_long2=pivot_longer(out,cols=all_of(c(n2:nend)),
                        names_to="subpop",values_to="infectious")
  
  # Plot results using ggplot2
  library(ggplot2)
  p1=ggplot(data=out_long1,aes(x=time,y=susceptible,group=subpop,colour=subpop)) +
    geom_line() +
    xlab("Time (days)") + ylab("Susceptible") +
    scale_color_manual(name = "",
                       labels=c("data1","data2","data3","data4","data5"),
                       values = c("green4","green3","green","limegreen","green2")) +
    scale_x_continuous(breaks=seq(0,60,10),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))
  
  p2=ggplot(data=out_long2,aes(x=time,y=infectious,group=subpop,colour=subpop)) +
    geom_line() +
    xlab("Time (days)") + ylab("Infectious") +
    scale_color_manual(name = "",
                       labels=c("data1","data2","data3","data4","data5"),
                       values = c("red4","red3","red","lightcoral","red2")) +
    #scale_y_continuous(expand=expansion(mult=c(0,0.05)),) +
    scale_x_continuous(breaks=seq(0,60,10),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