##
## NME COVID Model Tutorial: 
## Adding Asymptomatic Pathway and Screening-Based Interventions to COVID Model
##

library(EpiModel)

rm(list = ls())


# Network Initialization --------------------------------------------------

# Set age ranges and mortality rates by age
ages <- 0:85
departure_rate <- c(588.45, 24.8, 11.7, 14.55, 47.85, 88.2, 105.65, 127.2,
                    154.3, 206.5, 309.3, 495.1, 736.85, 1051.15, 1483.45,
                    2294.15, 3642.95, 6139.4, 13938.3)
dr_pp_pd <- departure_rate / 1e5 / 365
age_spans <- c(1, 4, rep(5, 16), 1)
dr_vec <- rep(dr_pp_pd, times = age_spans)

# Initialize network
n <- 1000
nw <- network_initialize(n)
ageVec <- sample(ages, n, replace = TRUE)
nw <- set_vertex_attribute(nw, "age", ageVec)

# Set nodal attributes
statusVec <- rep("s", n)
init.latent <- sample(1:n, 50)
statusVec[init.latent] <- "e"

statusTime <- rep(NA, n)
statusTime[which(statusVec == "e")] <- 1

clinical <- rep(NA, n)
dxStatus <- rep(0, n)

nw <- set_vertex_attribute(nw, "status", statusVec)
nw <- set_vertex_attribute(nw, "statusTime", statusTime)
nw <- set_vertex_attribute(nw, "clinical", clinical)
nw <- set_vertex_attribute(nw, "dxStatus", dxStatus)
nw


# Network Model Estimation ------------------------------------------------

# Parameterize network model
formation <- ~edges + degree(0) + absdiff("age")

mean_degree <- 2
edges <- mean_degree * (n/2)
avg.abs.age.diff <- 2
isolates <- n * 0.08
absdiff <- edges * avg.abs.age.diff

target.stats <- c(edges, isolates, absdiff)
target.stats

coef.diss <- dissolution_coefs(~offset(edges), 20, mean(dr_vec))
coef.diss

# Estimate network model
est <- netest(nw, formation, target.stats, coef.diss,
              set.control.ergm = control.ergm(MCMLE.maxit = 100))

# Model diagnostics
dx <- netdx(est, nsims = 10, ncores = 5, nsteps = 500,
            nwstats.formula = ~edges + absdiff("age") + degree(0:6),
            set.control.tergm = control.simulate.formula.tergm(MCMC.burnin.min = 20000))
print(dx)
plot(dx)


# Epidemic Model Simulation -----------------------------------------------

# Parameterize model
param <- param.net(inf.prob = 0.1,
                   act.rate = 3,
                   departure.rates = dr_vec,
                   departure.disease.mult = 1000,
                   arrival.rate = 1/(365*85),
                   inf.prob.a.rr = 0.5,
                   act.rate.dx.inter.time = Inf,
                   act.rate.dx.inter.rr = 0.05,
                   # proportion in clinical pathway by age decade
                   prop.clinical = c(0.40, 0.25, 0.37, 0.42, 0.51, 0.59, 0.72, 0.76),
                   ea.rate = 1/4.0,
                   ar.rate = 1/5.0,
                   eip.rate = 1/4.0,
                   ipic.rate = 1/1.5,
                   icr.rate = 1/3.5,
                   pcr.sens = 0.8,
                   dx.rate.sympt = 0.2,
                   dx.rate.other = 0.01,
                   allow.rescreen = FALSE)

# Initialize model
init <- init.net()

# Set control settings. Note this script runs a SINGLE simulation on one core so
# that it finishes quickly and can be stepped through in browser mode. The course
# page runs the same model with nsims = 10 on ncores = 5; raise these to match if
# you want to reproduce the interval bands in the plots on the website.
source("mod13-COVID2-fx.R")
control <- control.net(type = NULL,
                       nsims = 1,
                       ncores = 1,
                       nsteps = 100,
                       infection.FUN = infect2,
                       progress.FUN = progress2,
                       dx.FUN = dx_covid,
                       aging.FUN = aging,
                       departures.FUN = dfunc2,
                       arrivals.FUN = afunc2,
                       resimulate.network = TRUE,
                       tergmLite = TRUE,
                       set.control.tergm =
                         control.simulate.formula.tergm(MCMC.burnin.min = 10000))

# Simulate model
sim <- netsim(est, param, init, control)

sim


# Model Analysis ----------------------------------------------------------

# Plot of disease states over time
par(mfrow = c(1,2))
plot(sim)
plot(sim, y = c("a.num", "ip.num", "ic.num"), legend = TRUE)

# Plot of transitions
par(mfrow = c(1, 1))
plot(sim, y = c("se.flow", "ea.flow", "eip.flow"), legend = TRUE)

# Export mean data
df <- as.data.frame(sim, out = "mean")
head(df)

# Calculate average cumulative incidence
sum(df$se.flow, na.rm = TRUE)