# This is the R version of program 7.3 from page 256 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SIR epidemic on an nxn lattice with coupling (rho) to 4 nearest
# neighbour. nI defines the number of lattice sites (chosen randomly) that
# start with some infection.

# Set up equations as a function
CML = function(time,state,pars){
  with(as.list(c(state,pars)),{
    nsq=n*n
    X=state[1:nsq]; dXdt = rep(0,nsq)
    ystart=nsq+1; yend = 2*nsq;
    Y=state[ystart:yend]; dYdt = rep(0,nsq)
    
    N=matrix(1,nsq,1)
    
    dXdt = nu - mu*X - beta*X*(m%*%Y)/(m%*%N)
    dYdt = -gamma*Y - mu*Y + beta*X*(m%*%Y)/(m%*%N)
      
    return(list(c(dXdt,dYdt)))
  })
}

# Set parameters (all rates are per day)
n=25      # size of lattice, such that nxn subpopulations are arranged in 2D grid
beta=1.0  # Transmission rate within the subpopulation
gamma=0.1 # Recovery rate
mu=0.0001 # Per capita death rate
rho=0.1   # Rate at which individuals interact with their neighbouring environment

# Initial conditions
X0=0.1    # Initial number of susceptible individuals in a subpopulation
nI=4      # Number of subpopulations (randomly distributed) that contain infection
N0=1      # Initial number of individuals in a subpopulation
nu=N0*mu  # Total birth rate into each subpopulation
timestep=1
MaxTime=2910

# Check all the parameters are valid
if(X0<=0){
  print(paste("Initial number of susceptibles is less than or equal to zero, X0 =",X0))   
    stop()}
if(nI<=0){
  print(paste("Number of sources is less than or equal to zero, nI =",nI))  
  stop()}
if(N0<=0){
  print(paste("Initial number of individuals is less than or equal to zero, N0 =",N0))  
  stop()}
if(beta<=0){
  print(paste("Transmission rate is less than or equal to zero, beta =",beta))  
  stop()}
if(n<=0){
  print(paste("Lattice size is less than or equal to zero, n =",n))  
  stop()}
if(beta<=0){
  print(paste("Transmission rate is less than or equal to zero, beta =",beta))  
  stop()}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(mu<0){
  print(paste("Death rate is less than zero, mu =",mu))  
  stop()}
if(rho<0){
  print(paste("Interaction is less than zero, rho =",rho))  
  stop()}
if(rho>=0.25){
  print(paste("Interaction is greather than or equal to 1/4, rho =",rho))  
  stop()}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}

# Convert lattice to metapopulation
X = c(rep(X0,(n*n)))
Y = c(rep(0.0,(n*n)))
library(pracma)
Y[ceiling(n*n*rand(nI,1))] = 0.001*X0
V0 = c(X,Y)

m=((1-4*rho))*diag(1,(n*n))
ri=c(rep(seq(1,(n-1)),each=n))
rj=c(rep(1:n,(n-1)))
library(Matrix)
sm=sparseMatrix(ri*n+rj-n,(ri+1)*n+rj-n,x=rho, dims=c((n*n),(n*n)))
m=m+as.matrix(sm)
sm=sparseMatrix((ri+1)*n+rj-n,ri*n+rj-n,x=rho, dims=c((n*n),(n*n)))
m=m+as.matrix(sm)
sm=sparseMatrix(rj*n+ri-n,rj*n+ri+1-n,x=rho, dims=c((n*n),(n*n)))
m=m+as.matrix(sm)
sm=sparseMatrix(rj*n+(ri+1)-n,rj*n+ri-n,x=rho, dims=c((n*n),(n*n)))
m=m+as.matrix(sm)
pars = list(beta,gamma,mu,nu,N0,m,n)

# Initial state - reshape Y into nxn matrix
A0 = matrix(Y,nrow = n, ncol = n)
xx=c(); yy=c(); Variable=c()
for (i in 1:n){
  for (j in 1:n){
    xx=c(xx,i)
    yy=c(yy,j)
    Variable=c(Variable,A0[i,j])
  }}
Ai = data.frame(x=xx,y=yy,value=Variable)

# The main iteration
tt=rep(0,MaxTime)
Xs=rep(0,MaxTime)
Ys=rep(0,MaxTime)
t=0; tt[1]=0
Xs[1] = sum(X); Ys[1] = sum(Y)
i=2
nsq=n*n

times = seq(0,timestep,timestep)
library(deSolve)

for(i in 2:MaxTime){
  CML.out = ode(V0,times,CML,pars,method="ode45",rtol=1e-3)
  XX=tail(CML.out[,1:nsq+1],1)
  YY=tail(CML.out[,(nsq+2):(2*nsq+1)],1)
  V0=tail(CML.out[,2:(2*nsq+1)],1)
  Xs[i] = sum(XX); Ys[i] = sum(YY)
  tt[i] = tt[i-1]+timestep; 
  t=tt[i]
}

out = as.data.frame(cbind(tt,Xs,Ys))  

library(ggplot2)

#reshape YY into nxn matrix
A = matrix(YY,nrow = n, ncol = n)
xx=c(); yy=c(); Variable=c()
for (i in 1:n){
  for (j in 1:n){
    xx=c(xx,i)
    yy=c(yy,j)
    Variable=c(Variable,A[i,j])
  }}
AA = data.frame(x=xx,y=yy,value=Variable)

p1 = ggplot(Ai, aes(x=xx, y=yy, fill=value)) + 
  geom_tile(color = "black") +
  scale_fill_gradient(low = "red4", high = "red",name="Y(0)") +
  xlab("") + ylab("") +
  guides(fill = guide_colourbar(barwidth = 0.5, barheight = 5))+
  theme_classic() +
  theme(legend.title=element_text(size=9),panel.border = element_rect(color = "black",
                                    fill = NA, linewidth = 0.6))     
p2 = ggplot(AA, aes(x=xx, y=yy, fill=value)) + 
  geom_tile(color = "black") +
  scale_fill_gradient(low = "red4", high = "red",name="Y(MaxTime)") +
  xlab("") + ylab("") +
  guides(fill = guide_colourbar(barwidth = 0.5, barheight = 5))+
  theme_classic() +
  theme(legend.title=element_text(size=9),panel.border = element_rect(color = "black",
                                    fill = NA, linewidth = 0.6))  

p3=ggplot()+
  geom_line(data=out,aes(x=tt,y=Xs),linetype=1,color="green3")+
xlab("Time (days)") + ylab("Total Susceptibles") +
  scale_x_continuous(expand=expansion(mult=c(0,0.0)),breaks=c(0,500,1000,1500,2000,2500)) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black",
                fill = NA, linewidth = 0.6))          

p4=ggplot()+
  geom_line(data=out,aes(x=tt,y=Ys),linetype=1,color="red3")+
  xlab("Time (days)") + ylab("Total Infecteds") +
  scale_x_continuous(expand=expansion(mult=c(0,0.0)),breaks=c(0,500,1000,1500,2000,2500)) +
  theme_classic() +
  theme(panel.border = element_rect(color = "black",
                                    fill = NA, linewidth = 0.6))      

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