
infect <- function(dat, at) {

  ## Uncomment this to run environment interactively
  # browser()

  ## Attributes ##
  active <- get_attr(dat, "active")
  status <- get_attr(dat, "status")
  infTime <- get_attr(dat, "infTime")

  ## Parameters ##
  inf.prob <- get_param(dat, "inf.prob")
  act.rate <- get_param(dat, "act.rate")

  ## Find infected nodes ##
  idsInf <- which(active == 1 & status == "i")
  nActive <- sum(active == 1)
  nElig <- length(idsInf)

  ## Initialize default incidence at 0 ##
  nInf <- 0

  ## If any infected nodes, proceed with transmission ##
  if (nElig > 0 && nElig < nActive) {

    ## Look up discordant edgelist ##
    del <- discord_edgelist(dat, at)

    ## If any discordant pairs, proceed ##
    if (!(is.null(del))) {

      # Set parameters on discordant edgelist data frame
      del$transProb <- inf.prob
      del$actRate <- act.rate
      del$finalProb <- 1 - (1 - del$transProb)^del$actRate

      # Stochastic transmission process
      transmit <- rbinom(nrow(del), 1, del$finalProb)

      # Keep rows where transmission occurred
      del <- del[which(transmit == 1), ]

      # Look up new ids if any transmissions occurred
      idsNewInf <- unique(del$sus)
      nInf <- length(idsNewInf)

      # Set new attributes and transmission matrix
      if (nInf > 0) {
        status[idsNewInf] <- "e"
        infTime[idsNewInf] <- at
        dat <- set_attr(dat, "status", status)
        dat <- set_attr(dat, "infTime", infTime)

        dat <- set_transmat(dat, del, at)

      }
    }
  }

  ## Save summary statistic for S->E flow
  dat <- set_epi(dat, "se.flow", at, nInf)

  return(dat)
}


progress <- function(dat, at) {

  ## Uncomment this to function environment interactively
  # browser()

  ## Attributes ##
  active <- get_attr(dat, "active")
  status <- get_attr(dat, "status")

  ## Parameters ##
  ei.rate <- get_param(dat, "ei.rate")
  ir.rate <- get_param(dat, "ir.rate")

  ## E to I progression process ##
  nInf <- 0
  idsEligInf <- which(active == 1 & status == "e")
  nEligInf <- length(idsEligInf)

  if (nEligInf > 0) {
    vecInf <- which(rbinom(nEligInf, 1, ei.rate) == 1)
    if (length(vecInf) > 0) {
      idsInf <- idsEligInf[vecInf]
      nInf <- length(idsInf)
      status[idsInf] <- "i"
    }
  }

  ## I to R progression process ##
  nRec <- 0
  idsEligRec <- which(active == 1 & status == "i")
  nEligRec <- length(idsEligRec)

  if (nEligRec > 0) {
    vecRec <- which(rbinom(nEligRec, 1, ir.rate) == 1)
    if (length(vecRec) > 0) {
      idsRec <- idsEligRec[vecRec]
      nRec <- length(idsRec)
      status[idsRec] <- "r"
    }
  }

  ## Write out updated status attribute ##
  dat <- set_attr(dat, "status", status)

  ## Save summary statistics ##
  dat <- set_epi(dat, "ei.flow", at, nInf)
  dat <- set_epi(dat, "ir.flow", at, nRec)
  dat <- set_epi(dat, "e.num", at,
                 sum(active == 1 & status == "e"))
  dat <- set_epi(dat, "r.num", at,
                 sum(active == 1 & status == "r"))

  return(dat)
}




# LAB STEP 2 (part 1). The aging module now recalculates age.grp every step
# from the freshly incremented age, so a node crossing its 18th birthday
# moves group on its own. Recalculating here, rather than anywhere later, is
# what keeps age.grp current for the network resimulation that follows in
# the same step.
#
# Note a side effect of deriving age.grp wholesale from age like this: it is
# self-repairing. If the arrivals module forgot to append an age.grp value
# for a newborn, this line would rebuild the whole vector at full length on
# the next step and the gap would close on its own. That is a property of a
# DERIVED attribute, not a general rule. An attribute that cannot be
# recomputed from something else (a risk group assigned at entry, say) has
# no such safety net, and a missing append would leave its vector short
# against the node ids for the rest of the run, with no error raised.
#
# LAB STEP 2 (part 3). prop.minor is the summary statistic validating the
# new attribute: it starts near 18/86 (about 0.21, since ages are drawn
# uniformly from 0:85) and rises to roughly 0.23. The rise is driven mainly
# by adult-skewed mortality (age-graded death rates, multiplied 100-fold for
# infected nodes), not by the handful of age-0 arrivals. If it were flat at
# 0 or 1, or NA, the attribute is not being maintained. (mean() over the
# whole vector is right here because tergmLite = TRUE drops departed nodes
# from the attribute vectors.)
aging <- function(dat, at) {

  age <- get_attr(dat, "age")
  age <- age + 1/365
  dat <- set_attr(dat, "age", age)
  age.grp <- ifelse(age < 18, 0, 1)
  dat <- set_attr(dat, "age.grp", age.grp)

  dat <- set_epi(dat, "meanAge", at, mean(age, na.rm = TRUE))
  dat <- set_epi(dat, "prop.minor", at, mean(age.grp == 0, na.rm = TRUE))

  return(dat)
}


dfunc <- function(dat, at) {

  ## Attributes
  active <- get_attr(dat, "active")
  exitTime <- get_attr(dat, "exitTime")
  age <- get_attr(dat, "age")
  status <- get_attr(dat, "status")

  ## Parameters
  dep.rates <- get_param(dat, "departure.rates")
  dep.dis.mult <- get_param(dat, "departure.disease.mult")

  ## Query alive
  idsElig <- which(active == 1)
  nElig <- length(idsElig)

  ## Initialize trackers
  nDepts <- 0
  idsDepts <- NULL

  if (nElig > 0) {

    ## Calculate age-specific departure rates for each eligible node ##
    ## Everyone older than 85 gets the final mortality rate
    whole_ages_of_elig <- pmin(ceiling(age[idsElig]), 86)
    departure_rates_of_elig <- dep.rates[whole_ages_of_elig]

    ## Multiply departure rates for diseased persons
    idsElig.inf <- which(status[idsElig] == "i")
    departure_rates_of_elig[idsElig.inf] <- departure_rates_of_elig[idsElig.inf] * dep.dis.mult

    ## Simulate departure process
    vecDepts <- which(rbinom(nElig, 1, departure_rates_of_elig) == 1)
    idsDepts <- idsElig[vecDepts]
    nDepts <- length(idsDepts)

    ## Update nodal attributes
    if (nDepts > 0) {
      active[idsDepts] <- 0
      exitTime[idsDepts] <- at
    }
  }

  ## Set updated attributes
  dat <- set_attr(dat, "active", active)
  dat <- set_attr(dat, "exitTime", exitTime)

  ## Summary statistics ##
  dat <- set_epi(dat, "total.deaths", at, nDepts)

  # covid deaths
  covid.deaths <- length(intersect(idsDepts, which(status == "i")))
  dat <- set_epi(dat, "covid.deaths", at, covid.deaths)


  return(dat)
}


afunc <- function(dat, at) {

  ## Parameters ##
  n <- get_epi(dat, "num", at - 1)
  a.rate <- get_param(dat, "arrival.rate")

  ## Process ##
  nArrivalsExp <- n * a.rate
  nArrivals <- rpois(1, nArrivalsExp)

  # Update attributes
  if (nArrivals > 0) {
    dat <- append_core_attr(dat, at = at, n.new = nArrivals)
    dat <- append_attr(dat, "status", "s", nArrivals)
    dat <- append_attr(dat, "infTime", NA, nArrivals)

    dat <- append_attr(dat, "age", 0, nArrivals)
    dat <- append_attr(dat, "age.grp", 0, nArrivals)
  }

  ## Summary statistics ##
  dat <- set_epi(dat, "a.flow", at, nArrivals)

  return(dat)
}
