## ----setup, include=FALSE-----------------------------------------------------
knitr::opts_chunk$set(echo = TRUE)

library(hestia)
library(dplyr)
library(posterior)
library(ggplot2)


## -----------------------------------------------------------------------------

head(siir)


## -----------------------------------------------------------------------------

inf_process <- make_infection_model(
  transmit(from = "S",
           to = c("Is", "Ia"),
           source = c("Is", "Ia"),
           split = "phi"),
  progress(from = "Is", to = "R", gamma_s = NA),
  progress(from = "Ia", to = "R", gamma_a = NA),
  mult_ih_inf_probs = TRUE)



## -----------------------------------------------------------------------------
obs_process <- make_observation_model(
  pcr = c("S" = 0.05, "Is" = 0.95, "Ia" = 0.95, "R" = 0.05),
  igg = c("S" = 0.01, "Is" = 0.01, "Ia" = 0.01, "R" = 0.8),
  symp = c("S" = 0.03, "Is" = 1-1e-10, "Ia" = 0.03, "R" = 0.03))


## ----message=FALSE, warning=FALSE, eval=FALSE---------------------------------
# ## Inputs
# # 1. Infection process model
# # 2. Observation process model
# # 3. Raw data
# # 4. Covariates (optional)
# # 5. Starting state probabilities
# 
# siir_res <- run_model(inf_model = inf_process,
#                       obs_model = obs_process,
#                       data = siir,
#                       init_probs = c(1-3*1e-10, 1e-10, 1e-10, 1e-10),
#                       stan_opts = stan_options(iter = 1000))
# 
# 

## -----------------------------------------------------------------------------

# The output from the model in the previous code chunk is available in the
# siir_res package data object

# Get summary statistics. Parameters are returned on the natural (model) scale,
# so no back-transformation is needed.
draws_sum <- summarise_draws(siir_res, "mean", "median",
                             ~ quantile2(.x, probs = c(0.025, 0.975))) |>
  mutate(var_type = c(rep("Infection probability", 3),
                      rep("Recovery rate", 2),
                      "Symptomatic proportion"))

# "True" values from underlying simulation
draws_sum$true_value <- c(0.01, 0.05, 0.025, 0.2, 1/3, 0.7)

draws_sum


## ----fig.width=7, fig.height=5.5----------------------------------------------
draws_plot <- draws_sum |>
  mutate(
    category = factor(
      c(rep("Infection Probability (daily)", 3),
        rep("Recovery Rate (per day)", 2),
        "Symptomatic Proportion (split)"),
      levels = c("Infection Probability (daily)", "Recovery Rate (per day)", "Symptomatic Proportion (split)")
    ),
    label = factor(
      c("Extra-household prob (eh_prob)",
        "Intra-household symptomatic (ih_prob_Is)",
        "Intra-household asymptomatic (ih_prob_Ia)",
        "Symptomatic recovery (gamma_s)",
        "Asymptomatic recovery (gamma_a)",
        "Symptomatic proportion (phi)"),
      levels = rev(c("Extra-household prob (eh_prob)",
                 "Intra-household symptomatic (ih_prob_Is)",
                 "Intra-household asymptomatic (ih_prob_Ia)",
                 "Symptomatic recovery (gamma_s)",
                 "Asymptomatic recovery (gamma_a)",
                 "Symptomatic proportion (phi)"))
    )
  )

ggplot(draws_plot, aes(y = label)) +
  geom_errorbar(aes(xmin = q2.5, xmax = q97.5), width = 0, color = "#7570b3", linewidth = 1) +
  geom_point(aes(x = median, fill = "Posterior median"), size = 3, shape = 21, color = "#4d4785", stroke = 1.4) +
  geom_point(aes(x = true_value, fill = "True simulation value"), size = 3, shape = 23, color = "#a63603", stroke = 1.4) +
  facet_wrap(~category, scales = "free", ncol = 1) +
  scale_fill_manual(
    name = "",
    values = c("Posterior median" = "#ffffff", "True simulation value" = "#d95f02"),
    guide = guide_legend(override.aes = list(
      shape = c(21, 23),
      color = c("#4d4785", "#a63603"),
      fill = c("#ffffff", "#d95f02")
    ))
  ) +
  labs(
    title = "SIIR Model Parameter Estimates",
    subtitle = "Circle = posterior median; bar = 95% CrI; diamond = true simulation value",
    x = "Parameter Value",
    y = NULL
  ) +
  theme_minimal() +
  theme(
    legend.position = "top",
    legend.justification = "right",
    panel.grid.minor = element_blank(),
    panel.grid.major.y = element_line(color = "grey90"),
    plot.title = element_text(face = "bold"),
    strip.text = element_text(face = "bold", hjust = 0)
  )