##
## NME COVID Model Tutorial: Starting a COVID Model
##

library("EpiModel")

rm(list = ls())


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

# Range of possible ages
ages <- 0:85

# Age-specific mortality rates
# Rates per 100,000 for age groups: <1, 1-4, 5-9, 10-14, 15-19, 20-24, 25-29,
#                                   30-34, 35-39, 40-44, 45-49, 50-54, 55-59,
#                                   60-64, 65-69, 70-74, 75-79, 80-84, 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)

# Per-capita daily death rate (divide by 365,000)
dr_pp_pd <- departure_rate / 1e5 / 365

# Create vector of daily death rates
age_spans <- c(1, 4, rep(5, 16), 1)
dr_vec <- rep(dr_pp_pd, times = age_spans)
data.frame(ages, dr_vec)

# Plot death rates
par(mar = c(3,3,2,1), mgp = c(2,1,0), mfrow = c(1,1))
plot(ages, dr_vec, type = "o", xlab = "age", ylab = "Mortality Rate")

# Initialize network
n <- 1000
nw <- network_initialize(n)

# Set age attribute
ageVec <- sample(ages, n, replace = TRUE)
nw <- set_vertex_attribute(nw, "age", ageVec)

# Create vector of disease statuses
statusVec <- rep("s", n)
init.latent <- sample(1:n, 50)
statusVec[init.latent] <- "e"
which(statusVec == "e")
table(statusVec)

# Set status attribute
nw <- set_vertex_attribute(nw, "status", statusVec)
nw



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

# Review network object
nw

# Define the formation model: edges
formation <- ~edges + degree(0) + absdiff("age")

# Target statistics
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

# Dissolution model
coef.diss <- dissolution_coefs(~offset(edges), 20, mean(dr_vec)*2)
coef.diss

# Fit the model
est <- netest(nw, formation, target.stats, coef.diss)

# 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)


# EpiModel Model Simulation -----------------------------------------------

# Parameterize model
param <- param.net(inf.prob = 0.1,
                   act.rate = 3,
                   departure.rates = dr_vec,
                   departure.disease.mult = 100,
                   arrival.rate = 1/(365*85),
                   ei.rate = 0.05, ir.rate = 0.05)

# Initialize model
init <- init.net()

# Control settings (source function script)
source("mod13-COVID1-fx.R")
control <- control.net(type = NULL,
                       nsims = 5,
                       ncores = 5,
                       nsteps = 300,
                       infection.FUN = infect,
                       progress.FUN = progress,
                       aging.FUN = aging,
                       departures.FUN = dfunc,
                       arrivals.FUN = afunc,
                       resimulate.network = TRUE,
                       tergmLite = TRUE,
                       set.control.tergm =
                         control.simulate.formula.tergm(MCMC.burnin.min = 10000),
                       nwstats.formula = ~edges + meandeg + degree(0:2) + absdiff("age"))

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

# Show network statistics for fitted model terms
plot(sim, type = "formation",
     stats = c("edges", "absdiff.age", "degree0"), plots.joined = FALSE)

# Compare network stats for mean degree
plot(sim, type = "formation", stats = "meandeg", qnts = 1, ylim = c(1, 3))
abline(h = 2, lty = 2, lwd = 2)

# Simulation plot
plot(sim, qnts = 1)

# Disease incidence plotted 2 ways
par(mfrow = c(1,2))
plot(sim, y = "se.flow")
plot(sim, y = "se.flow", mean.smooth = FALSE, qnts = FALSE)

# Mean age summary statistic
par(mfrow = c(1,1))
plot(sim, y = "meanAge")

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

# Population size over time
plot(sim, y = "num")

# Deaths per day
plot(sim, y = c("total.deaths", "covid.deaths"), qnts = FALSE, legend = TRUE)

# Calculate sum of COVID deaths
sum(df$covid.deaths, na.rm = TRUE)

# Extract the transmission matrix for Simulation 1
tm1 <- get_transmat(sim)
head(tm1, 20)

# Plot the phylogram for the first seed
phylo.tm1 <- as.phylo.transmat(tm1)
plot(phylo.tm1[[1]], show.node.label = TRUE, root.edge = TRUE, cex = 0.75)