#USER INPUT------
save_directory <- 'put the pathname of the folder of this file here'
excel_file_name <- 'SEIR-Q3 Parameters.xlsx'

#LOADING LIBRARY--------
tryCatch({
  library(deSolve)
  library(ggplot2)
  library(gridExtra)
  library(reshape2)
}, error=function(cond){
  install.packages("deSolve")
  install.packages("ggplot2")
  install.packages("gridExtra")
  install.packages("reshape2")
  library(deSolve)
  library(ggplot2)
  library(gridExtra)
  library(reshape2)
})

#Processing---------
add_word <- function(x, word){
  return(paste(word, toString(x), sep=' '))
}
#Model Functions---------
SEIR <- function(t, X, pars){
  #parameters re-definition------
  L <- pars[1:n_sets]
  r <- pars[(n_sets+1):(2*n_sets)]
  B <- pars[(2*n_sets+1):(3*n_sets)]
  M <- pars[(3*n_sets+1):(4*n_sets)]
  Eps <- pars[(4*n_sets+1):(5*n_sets)]
  G <- pars[(5*n_sets+1):(6*n_sets)]
  p1 <- pars[(6*n_sets+1):(7*n_sets)]
  p2 <- pars[(7*n_sets+1):(8*n_sets)]
  p3 <- pars[(8*n_sets+1):(9*n_sets)]
  p4 <- pars[(9*n_sets+1):(10*n_sets)]
  D <- pars[(10*n_sets+1):(11*n_sets)]
  FP <- pars[(11*n_sets+1):(12*n_sets)]
  FN <- pars[(12*n_sets+1):(13*n_sets)]
  q <- pars[(13*n_sets+1):(14*n_sets)]
  Tr <- pars[(15*n_sets+1):(16*n_sets)]
  TD <- pars[(16*n_sets+1):(17*n_sets)]
  r2 <- pars[(17*n_sets+1):(18*n_sets)]
  D2 <- pars[(18*n_sets+1):(19*n_sets)]
  
  #r with threshold
  r[Tr<=t] <- r2[Tr<=t]
  
  #D with threshold
  D[TD<=t] <- D2[TD<=t]
  
  #calculate rates with parameters
  d1FP <- FP*(D*((1/p1)-0.5))^(-1) #false-positive at S sub-population
  d2 <- (1-FN)*(D*((1/p2)-0.5))^(-1) #diagnosed at E (false negatives excluded)
  d3 <- (1-FN)*(D*((1/p3)-0.5))^(-1) #diagnosed at I (false negatives excluded)
  d4FP <- FP*(D*((1/p4)-0.5))^(-1) #false-positive at R sub-population
  
  #state re-definition--------
  S <- X[1:n_sets]
  E <- X[(n_sets+1):(2*n_sets)]
  I <- X[(2*n_sets+1):(3*n_sets)]
  R <- X[(3*n_sets+1):(4*n_sets)]
  Qs <- X[(4*n_sets+1):(5*n_sets)]
  Q <- X[(5*n_sets+1):(6*n_sets)]
  Qr <- X[(6*n_sets+1):(7*n_sets)]
  I_T <- X[(7*n_sets+1):(8*n_sets)]
  
  #rate calculation-----------
  dS <- L - r*B*I*(1/N)*S - M*S - d1FP*S + q*Qs
  dE <- r*B*I*(1/N)*S - Eps*E - M*E - d2*E
  dI <- Eps*E - G*I - M*I - d3*I
  dR <- G*I - M*R - d4FP*R + q*Qr + q*Q
  
  dQs <- d1FP*S  - q*Qs
  dQ <- d2*E + d3*I - q*Q
  dQr <- d4FP*R - q*Qr
  
  dI_T <- r*B*I*(1/N)*S
  
  return(list(c(dS, dE, dI, dR, dQs, dQ, dQr, dI_T)))
}

#READ-----------
setwd(save_directory)
#initial state
init_val <- read_excel(excel_file_name, sheet=2)
init_val <- init_val[,2:length(init_val[1,])]
n_sets <- length(init_val[1,])
init_val <- matrix(unlist(init_val), ncol=length(init_val[1,]), byrow=F)
init_val <- rbind(init_val, replicate(n_sets, 0))

#parameter values
parameters <- read_excel(excel_file_name, sheet=1)[,2:(length(init_val[1,])+1)]
parameters <- matrix(unlist(parameters), ncol=length(parameters[1,]), byrow=F)

#PRE-CALCULATION-----------
#total population
N <- apply(init_val,2,sum)
#time
timepoints = seq (0, parameters[15,1], 0.1)

#linearize
parameters <- as.vector(t(parameters))
init_val <- as.vector(t(init_val))

#MAIN---------------
ode_SEIR <- ode(y=init_val, times=timepoints, func=SEIR, 
                parms=parameters, method='lsoda')

colnames(ode_SEIR) <- c('time',
                        'Susceptible', 'Susceptible', 'Susceptible',
                        'Exposed', 'Exposed', 'Exposed',
                        'Infected', 'Infected', 'Infected',
                        'Recovered', 'Recovered', 'Recovered',
                        'QuarantineS', 'QuarantineS', 'QuarantineS',
                        'QuarantineEI', 'QuarantineEI', 'QuarantineEI',
                        'QuarantineR', 'QuarantineR', 'QuarantineR',
                        'InfectedTot', 'InfectedTot', 'InfectedTot')

#PLOTS
#S---------
cur_interest <- cbind.data.frame(ode_SEIR[,1], ode_SEIR[,2:(n_sets+1)])
colnames(cur_interest) <- c('time', seq(1, (length(cur_interest)-1), 1))
cur_interest <- melt(cur_interest, id.vars='time')
cur_interest[,2] <- unlist(lapply(cur_interest[,2], add_word, word='Simulation'))
plt_S <- ggplot(data=cur_interest, aes(x=time, y=value, col=variable))+
  geom_line()+
  theme_bw()+
  theme(legend.position='none')+
  ylim(0, max(N))+
  xlab('time / days')+
  ylab('Number of People')+
  labs(title='Susceptible Population')

#E-------------------------
cur_interest <- cbind.data.frame(ode_SEIR[,1], ode_SEIR[,(1*n_sets+2):(2*n_sets+1)])
colnames(cur_interest) <- c('time', seq(1, (length(cur_interest)-1), 1))
cur_interest <- melt(cur_interest, id.vars='time')
cur_interest[,2] <- unlist(lapply(cur_interest[,2], add_word, word='Simulation'))
plt_E <- ggplot(data=cur_interest, aes(x=time, y=value, col=variable))+
  geom_line()+
  theme_bw()+
  theme(legend.position='none')+
  ylim(0, max(N))+
  xlab('time / days')+
  ylab('Number of People')+
  labs(title='Exposed Population')

#I-------------
cur_interest <- cbind.data.frame(ode_SEIR[,1], ode_SEIR[,(2*n_sets+2):(3*n_sets+1)])
colnames(cur_interest) <- c('time', seq(1, (length(cur_interest)-1), 1))
cur_interest <- melt(cur_interest, id.vars='time')
cur_interest[,2] <- unlist(lapply(cur_interest[,2], add_word, word='Simulation'))
plt_I <- ggplot(data=cur_interest, aes(x=time, y=value, col=variable))+
  geom_line()+
  theme_bw()+
  theme(legend.position='none')+
  ylim(0, max(N))+
  xlab('time / days')+
  ylab('Number of People')+
  labs(title='Infectious Population')

#R-------------
cur_interest <- cbind.data.frame(ode_SEIR[,1], ode_SEIR[,(3*n_sets+2):(4*n_sets+1)])
colnames(cur_interest) <- c('time', seq(1, (length(cur_interest)-1), 1))
cur_interest <- melt(cur_interest, id.vars='time')
cur_interest[,2] <- unlist(lapply(cur_interest[,2], add_word, word='Simulation'))
plt_R <- ggplot(data=cur_interest, aes(x=time, y=value, col=variable))+
  geom_line()+
  theme_bw()+
  theme(legend.position='none')+
  ylim(0, max(N))+
  xlab('time / days')+
  ylab('Number of People')+
  labs(title='Recovered Population')

#dS, dE, dI, dR, dQs, dQ, dQr, dI_T
#plots: Q_total  I_accumulative
#Q----------
cur_interest <- ode_SEIR[,(4*n_sets+2):(5*n_sets+1)]+ode_SEIR[,(5*n_sets+2):(6*n_sets+1)]+
  ode_SEIR[,(6*n_sets+2):(7*n_sets+1)]
cur_interest <- cbind.data.frame(ode_SEIR[,1], cur_interest)
colnames(cur_interest) <- c('time', seq(1, (length(cur_interest)-1), 1))
cur_interest <- melt(cur_interest, id.vars='time')
cur_interest[,2] <- unlist(lapply(cur_interest[,2], add_word, word='Simulation'))
plt_Q <- ggplot(data=cur_interest, aes(x=time, y=value, col=variable))+
  geom_line()+
  theme_bw()+
  theme(legend.position='none')+
  #ylim(0, max(N))+
  xlab('time / days')+
  ylab('log Number of People')+
  labs(title='Quarantined Population')

#Accumulative I------------
cur_interest <- cbind.data.frame(ode_SEIR[,1], ode_SEIR[,(7*n_sets+2):(8*n_sets+1)])
colnames(cur_interest) <- c('time', seq(1, (length(cur_interest)-1), 1))
cur_interest <- melt(cur_interest, id.vars='time')
cur_interest[,2] <- unlist(lapply(cur_interest[,2], add_word, word='Simulation'))
plt_IT <- ggplot(data=cur_interest, aes(x=time, y=value, col=variable))+
  geom_line()+
  theme_bw()+
  theme(legend.title=element_blank())+
  scale_y_continuous(trans='log10', limits=c(1, max(N)))+
  xlab('time / days')+
  ylab('Number of People')+
  labs(title='Acc. Number of Infection')
#Plotting
grid.arrange(plt_S, plt_E, plt_I, plt_R, plt_Q, plt_IT)