## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(collapse = TRUE, comment = "#>", fig.width = 7,
                      fig.height = 4.2, dpi = 110)

## ----setup, message = FALSE---------------------------------------------------
library(shapleydecomp)

## ----model--------------------------------------------------------------------
create_sir_model <- function(population = 1e5, years = 1990:2020) {
  structure(
    list(population = population, years = years,
         S0 = population - 10, I0 = 10, R0 = 0),
    class = "sir_model"
  )
}

create_sir_params <- function(years = 1990:2020) {
  yrs <- as.character(years)
  list(
    years            = years,
    base_beta        = 0.3,
    gamma            = 0.1,
    beta_reduction   = setNames(seq(0, 0.5,  length.out = length(years)), yrs),
    treatment_effect = setNames(seq(0, 0.15, length.out = length(years)), yrs)
  )
}

## ----methods------------------------------------------------------------------
run_simulation.sir_model <- function(model, params, ...) {
  n <- length(model$years)
  S <- I <- R <- inf <- deaths <- numeric(n)
  S[1] <- model$S0; I[1] <- model$I0; R[1] <- model$R0
  for (t in 2:n) {
    yr    <- as.character(model$years[t])
    beta  <- params$base_beta * (1 - params$beta_reduction[[yr]])
    gamma <- params$gamma + params$treatment_effect[[yr]]
    inf[t]    <- beta * S[t - 1] * I[t - 1] / model$population
    deaths[t] <- 0.01 * I[t - 1]
    rec       <- gamma * I[t - 1]
    S[t] <- S[t - 1] - inf[t]
    I[t] <- I[t - 1] + inf[t] - rec - deaths[t]
    R[t] <- R[t - 1] + rec
  }
  data.table::data.table(year = model$years, infections = inf, deaths = deaths)
}

summarize_results.sir_model <- function(model, sim_output, ...) sim_output

get_baseline_value.sir_model <- function(model, params, factor_name, cf_type,
                                         baseline_year, ...) {
  if (cf_type == CF_TOTAL_IMPACT) return(list(type = "zero"))
  list(type = "change_since",
       value = params[[factor_name]][[as.character(baseline_year)]])
}

build_counterfactual.sir_model <- function(model, params, factor_name,
                                           decomp_year, cf_type,
                                           baseline_values, ...) {
  idx <- params$years >= decomp_year
  params[[factor_name]][idx] <-
    if (cf_type == CF_TOTAL_IMPACT) 0 else baseline_values[[factor_name]]$value
  params
}

## ----register-----------------------------------------------------------------
.S3method("run_simulation",       "sir_model", run_simulation.sir_model)
.S3method("summarize_results",    "sir_model", summarize_results.sir_model)
.S3method("get_baseline_value",   "sir_model", get_baseline_value.sir_model)
.S3method("build_counterfactual", "sir_model", build_counterfactual.sir_model)

## ----validate-----------------------------------------------------------------
model  <- create_sir_model()
params <- create_sir_params()

validate_model_interface(
  model, params,
  factors       = c("beta_reduction", "treatment_effect"),
  summary_vars  = c("infections", "deaths"),
  baseline_year = 2000
)

## ----decompose----------------------------------------------------------------
res <- shapley_decompose(
  model, params,
  factors      = c("beta_reduction", "treatment_effect"),
  summary_vars = c("infections", "deaths"),
  decomp_years = 2000:2020,
  cf_type      = "total_impact",
  verbose      = FALSE
)
names(res)
head(res$decomp)

## ----totals-------------------------------------------------------------------
tot <- sequential_totals(res)
tot[measure == "infections" & year %in% c(2005, 2010, 2020)]

## ----plot-totals, fig.alt = "Infections averted by each intervention per year, stacked, 2000 to 2020."----
inf <- tot[measure == "infections"]
op <- par(mar = c(3, 4.5, 1, 1))
averted <- -t(as.matrix(inf[, .(beta_reduction, treatment_effect)]))
colnames(averted) <- inf$year
barplot(averted, col = c("#3b6ea8", "#d08c2b"), border = NA, las = 1,
        ylab = "Infections averted", ylim = c(0, max(colSums(averted)) * 1.1),
        legend.text = c("Prevention (beta_reduction)", "Treatment"),
        args.legend = list(x = "topleft", bty = "n"))
par(op)

## ----exact--------------------------------------------------------------------
off <- res$step_scenarios[step_year == 1999, .(year, infections)]
obs <- run_simulation(model, params)[, .(year, observed = infections)]
chk <- merge(merge(inf, obs, by = "year"), off, by = "year")
chk[, gap := beta_reduction + treatment_effect - (observed - infections)]
max(abs(chk$gap))

## ----change-since-------------------------------------------------------------
since <- shapley_decompose(
  model, params,
  factors       = c("beta_reduction", "treatment_effect"),
  summary_vars  = "infections",
  decomp_years  = 2000:2020,
  cf_type       = "change_since",
  baseline_year = 2000,
  verbose       = FALSE
)
sequential_totals(since)[year == 2020]