# This is the R version of program 3.4 from page 87 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SEIR model with four different age-groups and 
# yearly "movements" between the groups mimicking the school year

# S0, E0 and I0 are all vectors.

# Set up equations as a function
SEIR_4 = function(time,state,pars){
  with(as.list(c(state,pars)),{
    S=c(S1,S2,S3,S4); dSdt = c(0,0,0,0)
    E=c(E1,E2,E3,E4); dEdt = c(0,0,0,0)
    I=c(I1,I2,I3,I4); dIdt = c(0,0,0,0)
    X = (beta%*%I)
    for (a in 1:4){
      dSdt[a]=nu[a]*55/75 - X[a]*S[a] - mu[a]*S[a]
      dEdt[a]=X[a]*S[a] - sigma*E[a] - mu[a]*E[a]
      dIdt[a]=sigma*E[a] - gamma*I[a] - mu[a]*I[a]
    }
    return(list(c(dSdt,dEdt,dIdt)))
  })
}

# Set parameters (all rates are per day)
beta=matrix(c(2.089, 2.089, 2.086, 2.037,
              2.089, 9.336, 2.086, 2.037,
              2.086, 2.086, 2.086, 2.037,
              2.037, 2.037, 2.037, 2.037),
              nrow = 4, ncol = 4, byrow = TRUE)    # Matrix of transmission rates
sigma=1/8               # Rate at which individuals move from exposed to infectious class
gamma=1/5               # Recovery rate
nu=c(1/(55*365),0,0,0)  # Death rate in each age group (only adults die)
mu=c(0,0,0,1/(55*365))  # Birth rate into each class (only adults give birth)

n=c(6, 4, 10, 55)/75    # Proportion of population in each age class

pars = list(beta,sigma,gamma,nu,mu)

# Initial conditions
S0=c(0.05, 0.01, 0.01, 0.008)         # Initial proportions of population susceptible and in particular age class
E0=c(0.0001, 0.0001, 0.0001, 0.0001)  # Initial proportions of population exposed and in particular age class
I0=c(0.0001, 0.0001, 0.0001, 0.0001)  # Initial proportions of population infectious and in particular age class
y0=c(S1=S0[1],S2=S0[2],S3=S0[3],S4=S0[4],
     E1=E0[1],E2=E0[2],E3=E0[3],E4=E0[4],
     I1=I0[1],I2=I0[2],I3=I0[3],I4=I0[4])

MaxTime=100*365
step=1
T0=0

# Check all parameters are valid
for (a in 1:4) {
  if(I0[a]<0){
    print(paste("Initial level of infection in age group is less than zero, I0[",a,"] =",I0[a]))   
    stop()}
  if(E0[a]<0){
    print(paste("Initial level of exposed in age group is less than zero, E0[",a,"] =",E0[a]))   
    stop()}  
  if(S0[a]<0){
    print(paste("Initial level of susceptibles in age group is less than zero, S0[",a,"] =",S0[a]))   
    stop()} 
  if(S0[a]+E0[a]+I0[a]>n[a]){
    print(paste("Initial level of susceptibles+infecteds in age group is greater than group size S0+E0+I0[",a,"] =",S0[a]+E0[a]+I0[a], ", n[",a,"] =",n[a])) 
    stop()} 
}
if(length(which(beta<0))>0){
  print(paste("Transmission rate beta has a term that is less than zero"))
  stop()} 
if(sigma<=0){
  print(paste("Exposed to infectious rate is less than or equal to zero, sigma =",sigma))  
  stop()}
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}  
  
# Solve differential equations
i=0
while(T0<MaxTime)
{
  i=i+1
  times =seq(T0,T0+365,step)

  library(deSolve)
  SEIR_4.out = ode(y0,times,SEIR_4,pars,method="ode45",rtol=1e-5)

  out2 = as.data.frame(SEIR_4.out)  
  if(i==1){
    out=out2
  } else{
    out=rbind(out,out2) 
  }

# Yearly movement 
  Send=tail(SEIR_4.out,1)[2:5]
  Eend=tail(SEIR_4.out,1)[6:9]
  Iend=tail(SEIR_4.out,1)[10:13]
  
  S0[1]=Send[1]-Send[1]/6
  S0[2]=Send[2]+Send[1]/6 - Send[2]/4
  S0[3]=Send[3]+Send[2]/4 - Send[3]/10
  S0[4]=Send[4]+Send[3]/10
  
  E0[1]=Eend[1]-Eend[1]/6
  E0[2]=Eend[2]+Eend[1]/6 - Eend[2]/4
  E0[3]=Eend[3]+Eend[2]/4 - Eend[3]/10
  E0[4]=Eend[4]+Eend[3]/10
  
  I0[1]=Iend[1]-Iend[1]/6
  I0[2]=Iend[2]+Iend[1]/6 - Iend[2]/4
  I0[3]=Iend[3]+Iend[2]/4 - Iend[3]/10
  I0[4]=Iend[4]+Iend[3]/10
  
  y0=c(S1=S0[1],S2=S0[2],S3=S0[3],S4=S0[4],
       E1=E0[1],E2=E0[2],E3=E0[3],E4=E0[4],
       I1=I0[1],I2=I0[2],I3=I0[3],I4=I0[4])
  T0 = tail(SEIR_4.out,1)[1]
}

# Create long data frames for plotting
library(tidyr)
out_long1=pivot_longer(out,cols=c("S1","S2","S3","S4"),
                      names_to="age_group",values_to="susceptible")
out_long2=pivot_longer(out,cols=c("I1","I2","I3","I4"),
                       names_to="age_group",values_to="infectious")


# Plot results using ggplot2
library(ggplot2)
p1=ggplot(data=out_long1,aes(x=time/365,y=susceptible,group=age_group,colour=age_group)) +
    geom_line() +
    xlab("Time (years)") + ylab("Susceptible") +
    scale_color_manual(name = "",
                     labels=c("0-6","6-10","10-20","20+"),
                     values = c("green3","green4","green","springgreen")) +
    scale_x_continuous(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/365,y=infectious,group=age_group,colour=age_group)) +
    geom_line() +
    xlab("Time (years)") + ylab("Infectious") +
    scale_color_manual(name = "",
                     labels=c("0-6","6-10","10-20","20+"),
                     values = c("red4","red3","red2","lightcoral")) +
    scale_y_continuous(expand=expansion(mult=c(0,0.05)),trans="log10") +
    scale_x_continuous(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))
# Find proportion sero-positive
m=which(out[,1]>tail(out[,1],1)-365)
R=c(0,0,0,0)
for(i in 1:4){
  R[i]=1-mean(out[m,i+1])/n[i]
}
X=c(0, 0, 6, 6,  6, 6, 10, 10, 10, 10, 20, 20, 20, 20, 25, 25)
Y = c(0, R[1], R[1], 0, 0, R[2], R[2], 0, 0, R[3], R[3], 0, 0, R[4], R[4], 0)
fill.plot=data.frame(X,Y)
p3=ggplot(data=fill.plot,aes(x=X,y=Y)) +
     xlab("Age group")  + ylab("Proportion sero-positive") +
     geom_area(fill = "red",position = 'stack', colour="black") +
     scale_y_continuous(expand=expansion(mult=c(0,0.05))) +
     scale_x_continuous(expand=expansion(mult=c(0,0.0))) +
     theme_classic() +
     theme(panel.border = element_rect(color = "black",
                                  fill = NA, linewidth = 0.6))

library(patchwork)
p1/p2/p3