delta_expectation.Rmd

---
title: Net Reputational Incentive with Shocks 
output: 
  pdf_document:
    number_sections: yes
    fig_caption: yes
    keep_tex: no
    includes:
      in_header: takeup_workingpaper_header.sty
params:
  structural_takeup_version: 93
  reduced_form_takeup_version: 93
---

```{r setup, include=FALSE}
if (interactive()) {
  params = list(
    structural_takeup_version = 93,
    reduced_form_takeup_version = 93
  )
}
library(magrittr)
library(tidyverse)
library(latex2exp)
library(ggthemes)
library(cmdstanr)
library(tidybayes)

knitr::opts_chunk$set(echo = FALSE, warning = FALSE, cache = FALSE) 

source(file.path("multilvlr", "multilvlr_util.R"))
source("dist_structural_util.R")

ggplot2::theme_set(theme_minimal())
library(paletteer)
canva_palette_vibrant <- "Primary colors with a vibrant twist"
```

# Structural Model

## Adding non-prosocial shock

Previously, we modeled the take-up probability as 
\[ \bar{y}(b,c) = \prob{b - c + \mu\Delta^*(b,c) + V > 0}, \]
where 
\[\Delta^*(b,c) = \expect{V | b,c, y = 1} - \expect{V | b,c, y = 0}.\]
This was relatively simple and can be rewritten as
\[\Delta^*(b,c) = \Delta[v^*(b,c)] = \expect{V | V > v^*(b,c)} - \expect{V | V \leq v^*(b,c)},\]
where $v^*(b,c)$ is implicitly defined as
\[ v^*(b,c) = - b + c - \mu\Delta[v^*(c,d)]. \]

Now consider adding an unobservable shock $U$, and allowing $\mu$ to vary with cost,
\[ \bar{y}(b,c) = \prob{b - c + \mu(c)\Delta^*(b,c) + V + U > 0}. \]
We define $W = V + U$ and $w^*(b,c)$ is implicitly defined as
\[ w^*(b,c) = - b + c - \mu(c)\Delta[w^*(c,d)], \]
which is no longer the prosocial type cutoff. In calculating $\Delta^*$ individuals must now account for the non-prosocial shocks $U$, such that
\begin{equation}
\begin{aligned}
  \Delta^*(b,c) &= \Delta[w^*(b,c)] \\
           &=\expect{V | b,c, y = 1} - \expect{V | b,c, y = 0} \\ 
           &= \expect{V | W > w^*(b,c)} - \expect{V | W \leq w^*(b,c)} \\
           &= \int_{-\infty}^\infty v \prob{V = v | W > w^*(b,c)}\,\mathrm{d}v - \int_{-\infty}^\infty v \prob{V = v | W \leq w^*(b,c)}\,\mathrm{d}v \\
           &= \int_{-\infty}^\infty \frac{v\prob{W > w^*(b,c) | V = v}\prob{V = v}}{\prob{W > w^*(b,c)}}\,\mathrm{d}v - \int_{-\infty}^\infty \frac{v\prob{W \leq w^*(b,c) | V = v}\prob{V = v}}{\prob{W \leq w^*(b,c)}}\,\mathrm{d}v \\ 
           &= \int_{-\infty}^\infty \frac{v\prob{U > w^*(b,c) - v}\prob{V = v}}{\prob{W > w^*(b,c)}}\,\mathrm{d}v - \int_{-\infty}^\infty \frac{v\prob{U \leq w^*(b,c) - v}\prob{V = v}}{\prob{W \leq w^*(b,c)}}\,\mathrm{d}v \\ 
           &= \int_{-\infty}^\infty \frac{v \left[1 - F_u\left(w^*(b,c) - v\right)\right]f_v(v)}{1 - F_w\left(w^*(b,c)\right)}\,\mathrm{d}v - \int_{-\infty}^\infty\frac{v F_u\left(w^*(b,c) - v\right)f_v(v)}{F_w\left(w^*(b,c)\right)}\,\mathrm{d}v \\ 
           &= \frac{-1}{F_w\left(w^*(b,c)\right)[1 - F_w\left(w^*(b,c)\right)]} \int_{-\infty}^\infty v F_u\left(w^*(b,c) - v\right)f_v(v)\,\mathrm{d}v\label{eq:delta}.
\end{aligned}
\end{equation}
$f(\cdot)$ and $F(\cdot)$ are the pdf and cdf, respectively, for the subscript random variable (here $U, V,$ and $W$).

To get a sense of what flexibility this affords our model, consider how $\mu\Delta^*$ changes in two different scenarios: first, don't include the shock $u$ but allow $\mu$ to vary, second, hold $\mu$ fixed at 1 and allow the standard deviation of $u$ to vary. From here on I'm assuming that $V \sim N(0,1)$, $U \sim N(0, \sigma_u)$, and they are uncorrelated.

```{r, fig.width=8}
#| delta-theory
sim_int <- expand.grid(
  w = seq(-2, 2, 0.1),
  u_sd = seq(0.0, 2.0, 0.05),
  mu = seq(0.1, 1, 0.05)
) %>% 
  filter(u_sd == 0 | mu == 1) %>% 
  rowwise() %>% 
  mutate(
    mu_delta = mu * calculate_delta(w, sqrt(1 + u_sd^2), u_sd),
    mu_delta_deriv = mu * calculate_delta_deriv(w, sqrt(1 + u_sd^2), u_sd)
  ) %>% 
  ungroup() %>% 
  pivot_longer(c(mu_delta, mu_delta_deriv), names_to = "delta_type")

sim_int %>% 
  filter(delta_type == "mu_delta") %>%
  ggplot(aes(w, value)) +
  geom_line(
    aes(group = u_sd), 
    data = . %>% 
      filter(mu == 1) %>% 
      mutate(type = r"{Add shocks: Holding $\mu = 1$ and varying $\sigma_u$}")
  ) +
  geom_line(
    aes(group = mu), 
    data = . %>% 
      filter(u_sd == 0) %>% 
      mutate(type = r"{Original model: Holding $\sigma_u = 0$ and varying $\mu$}")
  ) + 
  geom_line(
    aes(group = mu), 
    color = "red",
    data = . %>% 
      filter(u_sd == 0, mu == 1) 
  ) + 
  labs(
    title = latex2exp::TeX(r"{$\mu\Delta\[w^*\]$ with and without shocks.}"),
    x = latex2exp::TeX("$w^*$"),
    y = "",
    caption = latex2exp::TeX(r"{Black lines move down as $\mu$ decreases or $\sigma_u$ increases. Red lines are with $\mu = 1$ and $\sigma_u = 0$.}")
  ) +
  facet_wrap(vars(type), labeller = as_labeller(latex2exp::TeX, default = label_parsed)) +
  NULL
```

\break

## The Complete Model

$$
\expect{Y^\textrm{takeup}(z,d)} = \expect{Y^\textrm{takeup} | Z = z, D = d} = \prob{\bar{b}(z,d) + \mu(z,d)\Delta^*(\bar{b}(z,d)) + U + V > 0}
$$

* Two types of shocks (as described above)
    \begin{align*}
      V &\sim \mathtt{Normal}(0,1) \\
      U &\sim \mathtt{Normal}(0,\sigma_u) \\
      W &= V + U \sim \mathtt{Normal}(0,\sigma_w) \\
      \sigma_u &\sim \mathtt{InverseGamma}(\alpha^\sigma_u, \beta^{\sigma}_u) \\
      \sigma_w &= \sqrt{1 + \sigma_u^2}
    \end{align*}
   
    - $V$ is the prosocial type.
    - $U$ is any other non-prosocial shocks.
    - $\alpha^\sigma_u$ and $\beta^\sigma_u$ are hyperparameters for $\sigma_u$ selected to keep it away from zero to avoid strong degeneracy in the model. This assumption only makes it harder to pick up the effect of social signaling.
    - We assume that $U$ and $V$ are independent.
* $\bar{b}(z,d)$ is the average net private benefit of deworming in response to incentive and distance treatments $z$ and $d$.
    \begin{align*}
      \bar{b}(z,d) &= z'\boldsymbol{\beta} + \delta \cdot d \\
      \\
      \boldsymbol{\beta} &= (\beta_\textrm{intercept}, \beta_\textrm{ink}, \beta_\textrm{calendar}, \beta_\textrm{bracelet})' \\
      \beta_\textrm{intercept} &\sim \mathtt{Normal}(0, \tau^\beta_\textrm{intercept}) \\ 
      \beta_\textrm{ink} &\sim \mathtt{Normal}(0, \tau^\beta_\textrm{ink}) \\ 
      \beta_\textrm{calendar} &\sim \mathtt{Normal}(0, \tau^\beta_\textrm{calendar}) \\ 
      \beta_\textrm{bracelet} &= \beta_\textrm{calendar} + \gamma\mu^\textrm{wtp} \\
      \gamma &\sim \mathtt{Normal}^+(0, \tau^\gamma)
    \end{align*}
    
    
    - $z$ is a binary vector representing assigned incentive treatment where the first cell is always 1.
    - The relationship between $\beta_\textrm{bracelet}$ and $\beta_\textrm{calendar}$ is calibrated using data from the WTP experiment. 
    - $\tau^\beta_\textrm{intercept}, \tau^\beta_\textrm{ink}, \tau^\beta_\textrm{calendar},$ and $\tau^\gamma$ are hyperparameters. 
    
* $\mu(z,d)$ parameterizes reputational concerns where
    \begin{align*}
      \mu(z,d) &= \mu_\circ \cdot \exp\left(\eta^\textrm{1ord}(z,d) - \eta^\textrm{1ord}(\textrm{control}, 0)\right) \\
      \\
      \mu_\circ &\sim \mathtt{Normal}^+(0, \tau^\mu) 
    \end{align*}
    
    - $\mu(z,d)$ are calibrated by the beliefs sub-model via the latent predictor $\eta^\textrm{1ord}(z,d)$.
    - Should the incentives and distance have no effect on visibility $\mu(z,d) = \mu_\circ$.
    - See Equation \eqref{eq:delta} on how $\Delta^*(\bar{b}(z,d))$ is defined.
    - $\tau^\mu$ is a hyperparameter.
    
* The WTP sub-model is
    \begin{align*}
      V^\textrm{wtp} &\sim \mathtt{Normal}(\mu^\textrm{wtp}, \sigma^\textrm{wtp}) \\
      \mu^\textrm{wtp} &\sim \mathtt{Normal}(0, \tau^{\mu,\textrm{wtp}}) \\
      \sigma^\textrm{wtp} &\sim \mathtt{Normal}(0, \tau^{\sigma,\textrm{wtp}})
    \end{align*}
    \begin{equation*}
    \prob{G = g, S(m) = s | \mu^\textrm{wtp}, \sigma^\textrm{wtp}} = 
      \begin{cases}
        \prob{V^{\textrm{wtp}} < -m} & s = 0 \land g = -1 \\ 
        \prob{V^{\textrm{wtp}} > m} & s = 0 \land g = 1 \\ 
        g\cdot\left(\prob{V^{\textrm{wtp}} < g \cdot m} - \prob{V^{\textrm{wtp}} < 0}\right) & s = 1
      \end{cases}
    \end{equation*}
    
    - $V^\textrm{wtp}$ is the difference in monetary valuation for bracelets vs. calendars.
    - This sub-model is estimated using observed choices in the WTP experiment:
        (a) Initial choice of incentive (calendar or bracelet), $G \in \{-1, 1\}$.
        (b) Decision to switch, $S(m) \in \{0, 1\}$, in response to a random price (Ksh), $m \in \{0, 10, 20,\ldots, 100\}$.
    - $\tau^{\mu,\textrm{wtp}}$ and $\tau^{\sigma,\textrm{wtp}}$ are hyperparameters.
    
* The beliefs sub-model is
    \begin{align*}
      Y^{\textrm{1ord}}(z,d) &\sim \mathtt{Binomial}(Y^{\textrm{rec}}, p^{\textrm{1ord}}) \\
      p^{\textrm{1ord}}(z,d) &= \mathsf{logit}^{-1}(\eta^\textrm{1ord}(z,d)) \\
      \eta^\textrm{1ord}(z,d) &= z'(\boldsymbol{\beta}^\textrm{1ord} + \boldsymbol{\delta}^\textrm{1ord}\cdot d) \\
      \beta^\textrm{1ord}_z &\sim \mathtt{Normal}(0, \tau^{\beta,\textrm{1ord}}_z) \\
      \delta^\textrm{1ord}_z &\sim \mathtt{Normal}(0, \tau^{\delta,\textrm{1ord}}_z)
    \end{align*}
    
    - $Y^\textrm{1ord}$ is the number of known community members, $Y^\textrm{rec}$, for whom a respondent reports knowing their deworming status.
    - $\tau^{\beta,\textrm{1ord}}$ and $\tau^{\beta,\textrm{1ord}}$ are hyperparameters.  
    
So given the parameters 

* $\boldsymbol{\theta}^\textrm{takeup} = \{\beta_\textrm{intercept}, \beta_\textrm{ink}, \beta_\textrm{calendar}, \gamma, \mu_\circ, \sigma_u, \sigma_w \}$.
* $\boldsymbol{\theta}^\textrm{wtp} = \{\mu^\textrm{wtp}, \sigma^\textrm{wtp}\}$.
* $\boldsymbol{\theta}^\textrm{1ord} = \{\beta^\textrm{1ord}_\textrm{intercept}, \beta^\textrm{1ord}_\textrm{ink}, \beta^\textrm{1ord}_\textrm{calendar},\beta^\textrm{1ord}_\textrm{bracelet},\delta^\textrm{1ord}_\textrm{intercept}, \delta^\textrm{1ord}_\textrm{ink}, \delta^\textrm{1ord}_\textrm{calendar},\delta^\textrm{1ord}_\textrm{bracelet}\}$.

we can compute the posterior

$$
\prob{\boldsymbol{\theta}^\textrm{takeup} | Y^\textrm{takeup}} \propto \prob{ Y^\textrm{takeup}| \boldsymbol{\theta}^\textrm{takeup}, \boldsymbol{\theta}^\textrm{wtp}, \boldsymbol{\theta}^\textrm{1ord}}\prob{ \boldsymbol{\theta}^\textrm{takeup}, \boldsymbol{\theta}^\textrm{wtp}, \boldsymbol{\theta}^\textrm{1ord}} 
$$
By using the WTP and beliefs data to calibrate the WTP and beliefs sub-models we can replace $\prob{ \boldsymbol{\theta}^\textrm{takeup}, \boldsymbol{\theta}^\textrm{wtp}, \boldsymbol{\theta}^\textrm{1ord}}$ with 

\begin{align*}
\prob{ \boldsymbol{\theta}^\textrm{takeup}, \boldsymbol{\theta}^\textrm{wtp}, \boldsymbol{\theta}^\textrm{1ord}|Y^\textrm{wtp},Y^\textrm{1ord},Y^\textrm{rec}} &= \prob{\boldsymbol{\theta}^\textrm{takeup}}\prob{\boldsymbol{\theta}^\textrm{wtp} | Y^\textrm{wtp}} \prob{\boldsymbol{\theta}^\textrm{1ord} | Y^\textrm{1ord},Y^\textrm{rec}} \\
&\propto \underbrace{\prob{\boldsymbol{\theta}^\textrm{takeup}}\prob{\boldsymbol{\theta}^\textrm{wtp}}\prob{\boldsymbol{\theta}^\textrm{1ord}}}_\textrm{priors}\prob{Y^\textrm{wtp} | \boldsymbol{\theta}^\textrm{wtp}} \prob{Y^\textrm{1ord},Y^\textrm{rec} | \boldsymbol{\theta}^\textrm{1ord}}  
\end{align*}

## Take-up Results

```{r load-processed-results}
dist_fit_data <- str_split(params$structural_takeup_version, ",") %>% 
  pluck(1) %>% 
  c(., 91) %>%
  as.integer() %>% 
  map_dfr(~ {
    temp_env <- new.env()
    load(file.path("temp-data", str_glue("processed_dist_fit{.x}.RData")), envir = temp_env)
    
    temp_env$dist_fit_data %>% 
      # filter(fct_match(model_type, "structural")) %>% 
      mutate(version = .x)
  }) %>% 
  group_by(model, fit_type, model_type) %>% 
  # # temp solution to get rid of a model I don't want
  filter(!(model == "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP" & version == 93)) %>%
  filter(model != "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HIER") %>%
  filter(version == 93 | (model == "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP" & version == 91)) %>%
  filter(min_rank(version) == n()) %>% 
  ungroup()




if (params$structural_takeup_version != params$reduced_form_takeup_version) {
  temp_env <- new.env()
  load(file.path("temp-data", str_glue("processed_dist_fit{params$reduced_form_takeup_version}.RData")), envir = temp_env)
  
  dist_fit_data %<>% 
    bind_rows(
      temp_env$dist_fit_data %>% 
        filter(fct_match(model_type, "reduced form")) %>% 
        mutate(version = as.integer(params$reduced_form_takeup_version))
    )
}



dist_fit_data %<>% 
  mutate(
    model_c = if_else(
      fit_type == "fit", 
      as.character(model), 
      str_c(model, "_PRIOR")
    ), 
    model = factor(model_c), 
    model_name = str_replace(model_c, "_PRIOR", "_WITHOUT_ANY_DATA"), 
    model_name = str_replace_all(model_name, "_", " "), 
    model_name = str_to_title(model_name),
    model_name = case_when(
      model == 'STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP' ~ "Structural Model P-Hat",
      model == 'STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HIGH_MU_WTP_VAL' ~ "Structural Model, High WTP MU",
      model == 'STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HIGH_SD_WTP_VAL' ~ "Structural Model, High WTP SD", 
      model == "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_NO_WTP_SUBMODEL" ~ "Struct p-hat, No WTP data",
      model == "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_NO_BELIEFS_SUBMODEL" ~ "Struct p-hat, No belief data",
      model == "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_NO_SUBMODELS" ~ "Struct p-hat, No submodel data",
      TRUE ~ model_name
    )
  )  %>%
  left_join(
    bind_rows(
      filter(., fct_match(model_type, "structural"), fct_match(fit_type, "fit")) %>%
        select(model, model_type, fit_type) %>% 
        # mutate(model_color = c(canva_pal(canva_palette_vibrant)(min(4, n())), canva_pal("Pop art")(max(0, n() - 4)))),
        mutate(
          model_color = paletteer_d("ggthemes::Tableau_20")[1:n()]
        ) %>%
      filter(., !fct_match(model_type, "structural") | !fct_match(fit_type, "fit")) %>%
        select(model, model_type, fit_type) %>% 
        mutate(model_color = if_else(fct_match(fit_type, "fit"), "black", "darkgrey"))
    ),
  by = c("model", "model_type", "fit_type"))


no_data_dist_fit_data = dist_fit_data %>%
  filter(str_detect(model, "_NO_"))

dist_fit_data = dist_fit_data %>%
  filter(!str_detect(model, "_NO_") | str_detect(model, "REDUCED"))

dist_fit_data = dist_fit_data %>%
  mutate(model_name = case_when(
    str_detect(model, "LOG_MU") ~ str_c(model_name, " Log Mu Rep"),
    str_detect(model, "LINEAR_MU") ~ str_c(model_name, " Linear Mu Rep"),
    TRUE ~ model_name
  ))

rf_analysis_data <- dist_fit_data %>% 
  filter(
    fct_match(model_type, "reduced form"),
    fct_match(fit_type, "fit"),
  ) %$% 
  stan_data[[1]]$analysis_data 
```

### Treatment Effects

Running a simple model (a non-hierarchical model with a linear cost), we get the following estimates.

```{r incentive-te, fig.width=8, fig.height=16}

fit_models = dist_fit_data %>%
  filter(fit_type != "prior-predict")  %>%
  pull(model) %>%
  as.character()

models_we_want = c(
  "STRUCTURAL_LINEAR_U_SHOCKS", 
  "STRUCTURAL_LINEAR_U_SHOCKS_NO_BELIEFS_DIST", 
  "STRUCTURAL_LINEAR_U_SHOCKS_PRIOR", 
  "STRUCTURAL_LINEAR_U_SHOCKS_NO_BELIEFS_DIST_PRIOR",
  "STRUCTURAL_LINEAR_U_SHOCKS_LOG_MU_REP",
  "STRUCTURAL_LINEAR_U_SHOCKS_LINEAR_MU_REP", 
  "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP", 
  "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HIGH_MU_WTP_VAL",
  "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HIGH_SD_WTP_VAL",
  "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HEIR_FOB", 
  "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP_HEIR_SOB", 
  "REDUCED_FORM_NO_RESTRICT", 
  "REDUCED_FORM_NO_RESTRICT_DIST_CTS", 
  fit_models
  )


dist_fit_data = dist_fit_data %>%
  mutate(
    model = factor(model, levels = c(levels(model), models_we_want) %>% unique())
  )

dist_fit_data %>%
  filter(
    (fct_match(model_type, "structural") & model %in% models_we_want) | 
      (fct_match(model_type, "reduced form") & fct_match(fit_type, "fit")),
  ) %>% 
  mutate(
    est_takeup_te =
      map_if(est_takeup_te, fct_match(model_type, "structural"),
             filter, mu_assigned_treatment_left == assigned_treatment_left, mu_assigned_treatment_right == assigned_treatment_right) %>%
        map(filter,
            (is.na(assigned_dist_group_left) & is.na(assigned_dist_group_right)) | (assigned_dist_group_left == assigned_dist_group_right),
            assigned_treatment_left != assigned_treatment_right,
            fct_match(assigned_treatment_right, c("control")),
            fct_match(assigned_treatment_left, "bracelet") | !fct_match(assigned_treatment_right, "calendar")),
    # model_color = canva_pal(canva_palette_vibrant)(n())
    model_color = paletteer_d("ggthemes::Tableau_20")[1:n()]
  ) %>% 
  select(model, model_name, est_takeup_te, fit_type, model_color) %>% 
  mutate(
    est_takeup_te = map(
      est_takeup_te,
      mutate,
      assigned_dist_group_left = fct_explicit_na(assigned_dist_group_left, "Combined") %>% 
        fct_relabel(str_to_title) %>% 
        fct_relevel("Combined"),
      assigned_treatment_left = str_to_title(assigned_treatment_left)
    )
  ) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_te, assigned_treatment_left) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_discrete("") +
        labs(
          title = "Incentive Average Treatment Effect"
        ) +
        ggforce::facet_col(vars(assigned_dist_group_left), 
                   space = "free",
                   scales = "free_y") +
        NULL,
      
      plot_estimand_hist(., est_takeup_te, iter_takeup_te, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.2)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group_left), rows = vars(assigned_treatment_left)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1, 
      rel_heights = c(1, 0.4)
    )
  } 

ggsave(
  "temp-plots/cts-rf-te.pdf", 
  width = 8, 
  height = 16
)



```


```{r signaling-te, fig.width=8, fig.height=13}
dist_fit_data %>% 
  filter(
    (fct_match(model_type, "structural") & fct_match(model, models_we_want))  
  ) %>% 
  select(model, model_name, est_takeup_te, fit_type, model_color) %>% 
  mutate(
    est_takeup_te = map(
      est_takeup_te,
      filter,
      (is.na(assigned_dist_group_left) & is.na(assigned_dist_group_right)) | (assigned_dist_group_left == assigned_dist_group_right),
      across(c(assigned_treatment_left, assigned_treatment_right), fct_match, "control"),
      !is.na(mu_assigned_treatment_left),
      fct_match(mu_assigned_treatment_left, "bracelet") | !fct_match(mu_assigned_treatment_right, "calendar"),
      fct_match(mu_assigned_treatment_right, "control"),
    ) %>% 
      map(
        mutate,
        assigned_dist_group_left = fct_explicit_na(assigned_dist_group_left, "Combined") %>% 
          fct_relabel(str_to_title) %>% 
          fct_relevel("Combined"),
        mu_assigned_treatment_left = str_to_title(mu_assigned_treatment_left),
      ),
    # model_color = canva_pal(canva_palette_vibrant)(n())
    model_color = paletteer_d("ggthemes::Tableau_20")[1:n()]
  ) %>% { 
    cowplot::plot_grid(
      plot_estimands(., est_takeup_te, mu_assigned_treatment_left) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.05)) +
        scale_y_discrete("") +
        labs(
          title = "Signaling Average Treatment Effect",
          subtitle = "Holding private incentive at the control level.") +
        ggforce::facet_col(vars(assigned_dist_group_left), 
                   space = "free",
                   scales = "free_y") +
        NULL,
      
      plot_estimand_hist(., est_takeup_te, iter_takeup_te, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.2)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group_left), rows = vars(mu_assigned_treatment_left)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1, rel_heights = c(1, 0.4)
    )
  } 

ggsave(
  "temp-plots/anne-check/all-signalling.png", 
  width = 10, 
  height = 15,
  dpi = 600
)
```

```{r private-incentive-te, fig.width=8, fig.height=15}
dist_fit_data %>% 
  filter(
    (fct_match(model_type, "structural") & fct_match(model, models_we_want))  
  ) %>% 
  select(model, model_name, est_takeup_te, fit_type, model_color) %>% 
  mutate(
    est_takeup_te = map(
      est_takeup_te,
      filter,
      (is.na(assigned_dist_group_left) & is.na(assigned_dist_group_right)) | (assigned_dist_group_left == assigned_dist_group_right),
      !is.na(mu_assigned_treatment_left),
      !is.na(mu_assigned_treatment_right),
      across(c(mu_assigned_treatment_left, mu_assigned_treatment_right), fct_match, "control"),
      fct_match(assigned_treatment_left, "bracelet") | !fct_match(assigned_treatment_right, "calendar"),
      fct_match(assigned_treatment_right, "control"),
    ) %>% 
      map(
        mutate,
        assigned_dist_group_left = fct_explicit_na(assigned_dist_group_left, "Combined") %>% 
          fct_relabel(str_to_title) %>% 
          fct_relevel("Combined"),
        assigned_treatment_left = str_to_title(assigned_treatment_left),
      ),
    # model_color = canva_pal(canva_palette_vibrant)(n())
    model_color = paletteer_d("ggthemes::Tableau_20")[1:n()]
  ) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_te, assigned_treatment_left) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.05)) +
        scale_y_discrete("") +
        labs(
          title = "Private Incentive Average Treatment Effect",
          subtitle = "Holding signaling at the control level.") +
        ggforce::facet_col(vars(assigned_dist_group_left), space = "free", scales = "free_y") +
        NULL,
      
      plot_estimand_hist(., est_takeup_te, iter_takeup_te, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group_left), rows = vars(assigned_treatment_left)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1, rel_heights = c(1, 0.4)
    )
  }
ggsave(
  "temp-plots/anne-check/all-private.png", 
  width = 10, 
  height = 15,
  dpi = 600
)
```

```{r take-up-prob, fig.width=9}

dist_prob_plot <- dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want), fct_match(model_type, "structural")) %>% #, "STRUCTURAL_LINEAR_U_SHOCKS_NO_SUBMODELS"))) %>%
  transmute(
    model_name,
    cluster_takeup_prop = map(cluster_takeup_prop, ~ mutate(.x, roc_distance = roc_distance / 1000)),
    obs_cluster_takeup_level,
  ) %>% 
  ggplot(aes(roc_distance)) +
  geom_line(aes(y = per_0.5, color = model_name), data = . %>%  unnest(cluster_takeup_prop)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75, fill = model_name), alpha = 0.3, data = . %>%  unnest(cluster_takeup_prop)) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9, fill = model_name), alpha = 0.3, data = . %>%  unnest(cluster_takeup_prop)) +
  geom_rug(
    aes(dist),
    alpha = 0.5,
    data = rf_analysis_data %>%
      distinct(cluster_id, assigned.treatment, dist = cluster.dist.to.pot / 1000) %>% 
      rename(assigned_treatment = assigned.treatment) 
  ) +
  labs(
    x = "Distance to Treatment (d) [km]", y = latex2exp::TeX("Take-up Probability, $1 - F_w(w^*(z,d))$") ,
    caption = "Line: Median. Outer ribbon: 80% credible interval. Inner ribbon: 50% credible interval." 
      
  ) +
  coord_cartesian(xlim = c(0, 2.5), ylim = c(0, 0.55)) +
  NULL

dist_prob_plot +
  geom_linerange(aes(assigned_dist_obs / 1000, ymin = per_0.1, ymax = per_0.9, color = model_name), data = . %>% unnest(obs_cluster_takeup_level)) +
  geom_point(aes(assigned_dist_obs / 1000, y = obs_prop_takeup), size = 0.25, show.legend = FALSE, data = . %>% unnest(obs_cluster_takeup_level)) +
  scale_color_discrete("", aesthetics = c("color", "fill")) +
  facet_wrap(vars(assigned_treatment), labeller = labeller(assigned_treatment = str_to_title)) +
  # facet_grid(rows = vars(assigned_treatment), cols = vars(model_name), labeller = labeller(.default = str_to_title)) +
  theme(legend.position = "top") +
  NULL
```

### Rates of Change

To calculate the slope of $\bar{y}(b,c)$ with respect to cost,
\begin{align*}
  \frac{\partial \bar{y}(b,c)}{\partial c} = \frac{-f_w(w^*(b,c))\cdot(1 - \mu'(c)\Delta[w^*(b,c)])}{1 + \mu\Delta'[w^*(b,c)]},
\end{align*}
where
\begin{align*}
\Delta'[w] = &-\frac{\int_{-\infty}^\infty v f_u(w - v) f_v(v)\,\mathrm{d}v + f_w(w)\left[1 - 2F_w(w)\right]\Delta[w]}{F_w(w)\left[1 - F_w(w)\right]} 
\end{align*}

```{r rate-of-change}
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want), fct_match(model_type, "structural")) %>%
  transmute(
    model_name,
    cluster_roc = map2(cluster_roc, stan_data, ~ { 
      mutate(.x,
        roc_distance = roc_distance / 1000,
        across(starts_with("per_"), divide_by, sd(.y$analysis_data$cluster.dist.to.pot)), 
        across(starts_with("per_"), multiply_by, 1000),
        across(starts_with("per_"), multiply_by, 100)
      ) 
    })
  ) %>%
  unnest(cluster_roc) %>% 
  filter(fct_match(assigned_treatment, c("control", "ink", "bracelet"))) %>% 
  filter(roc_distance < 2.5) %>%
  ggplot(aes(roc_distance)) +
  geom_line(aes(y = per_0.5, color = model_name)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75, fill = model_name), alpha = 0.4) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9, fill = model_name), alpha = 0.4) +
  geom_rug(
    aes(dist),
    alpha = 0.75,
    data = rf_analysis_data %>%
      filter(fct_match(assigned.treatment, c("control", "ink", "bracelet"))) %>%
      distinct(cluster_id, assigned_treatment = assigned.treatment, dist = cluster.dist.to.pot / 1000)) +
  # scale_color_canva("Treatment (z)", labels = str_to_title, palette = "Fun and professional") +
  # scale_fill_canva("Treatment (z)", labels = str_to_title, palette = "Fun and professional") +
  scale_color_discrete("", aesthetics = c("color", "fill")) +
  labs(
    title = "Rates of Change",
    x = "Distance to Treatment (d) [km]", y = latex2exp::TeX(r"{Rate of Change \[pp/km\]}") ,
    caption = "Line: Median. Outer ribbon: 80% credible interval. Inner ribbon: 50% credible interval." 
      
  ) +
  facet_wrap(vars(assigned_treatment), labeller = labeller(.default = str_to_title)) +
  theme(legend.position = "top") +
  NULL
```

```{r}
#| rate-of-change-difference
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want), fct_match(model_type, "structural")) %>%
  transmute(
    model, model_name,
    y_rate_of_change_diff = map2(cluster_roc_diff, stan_data, ~ {
      mutate(.x,
        roc_distance = roc_distance / 1000,
        across(starts_with("per_"), divide_by, sd(.y$analysis_data$cluster.dist.to.pot)), 
        across(starts_with("per_"), multiply_by, 1000),
        across(starts_with("per_"), multiply_by, 100)
      ) 
  })) %>%
  unnest(y_rate_of_change_diff) %>% 
  filter(fct_match(assigned_treatment, c("ink", "bracelet"))) %>% 
  filter(roc_distance < 2.5) %>%
  ggplot(aes(roc_distance)) +
  geom_line(aes(y = per_0.5, color = model_name)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75, fill = model_name), alpha = 0.25) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9, fill = model_name), alpha = 0.25) +
  scale_color_discrete("", aesthetics = c("color", "fill")) +
  labs(
    title = "Difference in Rate of Change",
    x = "Distance to Treatment [km]", y = "Difference in Rate of Change [pp/km]",
    caption = "Line: Median. Outer ribbon: 80% credible interval. Inner ribbon: 50% credible interval." 
  ) +
  facet_wrap(vars(assigned_treatment), labeller = labeller(.default = str_to_title)) +
  theme(legend.position = "top") +
  # coord_cartesian(ylim = c(0, 0.15)) + 
  NULL
```

```{r, fig.height=2, fig.width=8}
#| difference-in-roc-difference
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want)) %>% 
  select(model_name, cluster_roc_diff_diffdist) %>% 
  unnest(cluster_roc_diff_diffdist) %>% 
  filter(fct_match(assigned_treatment, c("ink", "control", "bracelet"))) %>% 
  ggplot() +
  geom_pointinterval(aes(y = assigned_treatment, x = per_0.5, xmin = per_0.1, xmax = per_0.9, color = model_name), position = "dodge") +
  geom_vline(xintercept = 0, linetype = "dotted") +
  scale_x_continuous("Difference in Difference") +
  scale_y_discrete("", labels = str_to_title) +
  scale_color_discrete("", aesthetics = c("color", "fill")) +
  labs(title = "Difference in Rate of Change at 0 km and 2.5 km", caption = "Line interval shows the 80% credible interval.") +
  theme(legend.position = "top") +
  NULL
```

```{r}
#| type-cutoff
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want), fct_match(model_type, "structural" )) %>% 
  transmute(model_name, cluster_w_cutoff = map(cluster_w_cutoff, mutate, roc_distance = roc_distance / 1000)) %>% 
  unnest(cluster_w_cutoff) %>% 
  filter(fct_match(assigned_treatment, c("ink", "control", "bracelet"))) %>% 
  filter(roc_distance < 2.5) %>%
  ggplot(aes(roc_distance)) +
  geom_line(aes(y = per_0.5, color = assigned_treatment)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75, fill = assigned_treatment), alpha = 0.4) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9, fill = assigned_treatment), alpha = 0.4) +
  scale_color_discrete("", aesthetics = c("color", "fill")) +
  labs(
    title = "Type Cutoff",
    x = "Distance to Treatment", y = ""
  ) +
  # coord_cartesian(ylim = c(0, 0.1))
  # facet_wrap(vars(assigned_treatment), labeller = labeller(.default = str_to_title)) +
  theme(legend.position = "top")
  NULL


ggsave(
  "temp-plots/anne-check/vstart.png", 
  width = 10, 
  height = 10,
  dpi = 500
)
```

```{r, eval=FALSE}
# dist_fit_data %>% 
#   filter(fct_match(fit_type, "fit"), fct_match(model, c("STRUCTURAL_LINEAR_U_SHOCKS"))) %>% #, "STRUCTURAL_LINEAR_U_SHOCKS_NO_SUBMODELS"))) %>%
#   transmute(cluster_social_multiplier = map(cluster_social_multiplier, ~ mutate(.x, roc_distance = roc_distance / 1000))) %>% 
#   unnest(cluster_social_multiplier) %>% 
#   filter(fct_match(assigned_treatment, c("ink", "control", "bracelet"))) %>% 
#   ggplot(aes(roc_distance)) +
#   geom_line(aes(y = per_0.5, color = assigned_treatment)) +
#   # geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75, fill = treatment), alpha = 0.3) +
#   # geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9, fill = treatment), alpha = 0.3) +
#   labs(
#     x = "Distance to Treatment", y = ""
#   ) +
#   # coord_cartesian(ylim = c(0, 0.1))
#   NULL
```

```{r}
#| sim-delta, eval = FALSE
dist_fit_data %>% 
  filter(fct_match(model, models_we_want)) %>% 
  select(model_name, fit_type, sim_delta) %>% 
  unnest(sim_delta) %>% 
  ggplot(aes(w, fill = fit_type, group = fit_type)) +
  geom_line(aes(y = per_0.5)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75), alpha = 0.3) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9), alpha = 0.3) +
  labs(
    x = "w", y = ""
  ) +
  # coord_cartesian(ylim = c(0, 0.1))
  facet_wrap(vars(model_name)) +
  NULL

dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want)) %>% 
  select(model_name, sim_delta) %>% 
  unnest(sim_delta) %>% 
  ggplot(aes(w)) +
  geom_line(aes(y = per_0.5)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75), alpha = 0.3) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9), alpha = 0.3) +
  labs(
    x = "w", y = ""
  ) +
  # coord_cartesian(ylim = c(0, 0.1))
  facet_wrap(vars(model_name)) +
  NULL
```

```{r, eval=FALSE}
# dist_fit_data %>% 
#   filter(fct_match(fit_type, "fit"), fct_match(model, c("STRUCTURAL_LINEAR_U_SHOCKS"))) %>% #, "STRUCTURAL_LINEAR_U_SHOCKS_NO_SUBMODELS"))) %>%
#   pluck("y_rate_of_change_diff", 1) %>% 
#   mutate(
#     roc_distance = roc_distance / 1000,
#     across(starts_with("per_"), multiply_by, 1000)
#   ) %>% 
#   ggplot(aes(roc_distance)) +
#   geom_line(aes(y = per_0.5)) +
#   geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75), alpha = 0.25) +
#   geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9), alpha = 0.25) +
#   labs(
#     x = "Distance to Treatment", y = ""
#   ) 
```

### Valuation of Reputational Returns in Distance

Here estimate what is the distance, $r^*(z,d)$, that would have a disutility equal to the reputational returns at different $(z,d)$.

\begin{align*}
r^*(z,d)\delta &= \mu(z,d) \Delta[w^*(z,d)] \\
r^*(z,d) &= \frac{\mu(z,d) \Delta[w^*(z,d)]}{\delta} \\
\end{align*}

```{r, fig.width=8}
#| value-rep-returns
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want)) %>%
  select(model_name, cluster_rep_return_dist) %>%
  unnest(cluster_rep_return_dist) %>% 
  mutate(
    roc_distance = roc_distance / 1000,
    across(starts_with("per_"), divide_by, 1000)
  ) %>% 
  filter(roc_distance < 2.5) %>%
  ggplot(aes(roc_distance)) +
  geom_line(aes(y = per_0.5, color = model_name)) +
  geom_ribbon(aes(ymin = per_0.25, ymax = per_0.75, fill = model_name), alpha = 0.4) +
  geom_ribbon(aes(ymin = per_0.1, ymax = per_0.9, fill = model_name), alpha = 0.4) +
  scale_color_discrete("", aesthetics = c("color", "fill")) +
  labs(
    title = "Valuation of Reputational Returns in Terms of Distance",
    x = "Distance to Treatment [km]", y = "Distance Value [km]",
  ) +
  facet_wrap(vars(assigned_treatment), labeller = labeller(.default = str_to_title), nrow = 1) +
  theme(legend.position = "top") +
  NULL
```

### Levels

```{r, fig.width=8, fig.height=19}
#| levels
dist_fit_data %>% 
  filter(
    (fct_match(model_type, "structural") & fct_match(model, models_we_want)) | 
      (fct_match(model_type, "reduced form") & fct_match(fit_type, "fit")),
  ) %>% 
  select(model, model_type, model_name, est_takeup_level, fit_type, model_color) %>% 
  mutate(
    est_takeup_level = map_if(
      est_takeup_level, 
      !fct_match(model_type, "reduced form"),
      filter, assigned_treatment == mu_assigned_treatment
    ) %>%  
      map(
      mutate,
      assigned_dist_group = fct_explicit_na(assigned_dist_group, "Combined") %>% 
        fct_relabel(str_to_title) %>% 
        fct_relevel("Combined"),
      assigned_treatment = str_to_title(assigned_treatment),
    ),
    # model_color = canva_pal(canva_palette_vibrant)(n())
    model_color = paletteer_d("ggthemes::Tableau_20")[1:n()]
  ) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_level, assigned_treatment) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_discrete("") +
        labs(
          title = "Incentive Take-up Level"
        ) +
        ggforce::facet_col(vars(assigned_dist_group), space = "free", scales = "free_y") +
        NULL,
      
      plot_estimand_hist(., est_takeup_level, iter_prop_takeup, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group), rows = vars(assigned_treatment), scales = "free_y") +
        theme(legend.position = "none") +
        coord_cartesian(ylim = c(0, 0.2)) +
        NULL,
      
      ncol = 1, rel_heights = c(1, 0.4)
    )
  }
```

```{r, fig.width=8, fig.height=19}
#| signal-levels
dist_fit_data %>% 
  filter(
    (fct_match(model_type, "structural") & fct_match(model, models_we_want))  
  ) %>% 
  select(model, model_name, model_type, est_takeup_level, fit_type, model_color) %>% 
  mutate(
    est_takeup_level = map_if(
      est_takeup_level, 
      !fct_match(model_type, "reduced form"),
      filter,
      fct_match(assigned_treatment, "control"),
      !is.na(mu_assigned_treatment),
    ) %>% 
      map(
        mutate,
        assigned_dist_group = fct_explicit_na(assigned_dist_group, "Combined") %>% 
          fct_relabel(str_to_title) %>% 
          fct_relevel("Combined"),
        
        mu_assigned_treatment = str_to_title(mu_assigned_treatment),
      ),
    model_color = canva_pal(canva_palette_vibrant)(n())
  ) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_level, mu_assigned_treatment) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_discrete("", labels = str_to_title) +
        labs(
          title = "Signaling Level",
          subtitle = "Holding private incentive at the control level."
        ) +
        ggforce::facet_col(vars(assigned_dist_group), space = "free", scales = "free_y") +
        NULL,
      
      plot_estimand_hist(., est_takeup_level, iter_prop_takeup, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group), rows = vars(mu_assigned_treatment)) +
        coord_cartesian(ylim = c(0, 0.2)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1, rel_heights = c(1, 0.4)
    )
  }
  
```

```{r, fig.width=8, fig.height=19}
#| private-levels
dist_fit_data %>% 
  filter(
    (fct_match(model_type, "structural") & fct_match(model, models_we_want))  
  ) %>% 
  select(model, model_name, est_takeup_level, fit_type, model_color) %>% 
  mutate(
    est_takeup_level = map(
      est_takeup_level,
      filter,
      !is.na(mu_assigned_treatment),
      fct_match(mu_assigned_treatment, "control"),
    ) %>% 
      map(
        mutate,
        assigned_dist_group = fct_explicit_na(assigned_dist_group, "Combined") %>% 
          fct_relabel(str_to_title) %>% 
          fct_relevel("Combined"),
        
        assigned_treatment = str_to_title(assigned_treatment)
      ),
    model_color = canva_pal(canva_palette_vibrant)(n())
  ) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_level, assigned_treatment) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_discrete("") +
        labs(
          title = "Private Incentive Level",
          subtitle = "Holding signaling at the control level."
        ) +
        ggforce::facet_col(vars(assigned_dist_group), space = "free", scales = "free_y") +
        NULL,
      
      plot_estimand_hist(., est_takeup_level, iter_prop_takeup, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group), rows = vars(assigned_treatment)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1, rel_heights = c(1, 0.4)
    )
  }
```

<!-- ## Model Parameters -->

```{r, eval=FALSE}
dist_fit_data %>% 
  filter(fct_match(model_type, "structural")) %>% 
  select(fit_type, obs_cluster_mu_rep) %>% 
  unnest(obs_cluster_mu_rep) %>% 
  # filter(str_detect(variable, "mu_beliefs_effect")) %>% 
  filter(str_detect(variable, "obs_cluster_mu_rep")) %>% 
  unnest(iter_data) %>% 
  ggplot() +
  # geom_histogram(aes(iter_est, y = stat(density) * 0.05, fill = fit_type), 
  #                binwidth = 0.05, boundary = 0, position = "identity", alpha = 0.5) +
  geom_histogram(aes(iter_est), 
                 binwidth = 0.5, boundary = 0, position = "identity", alpha = 0.5) +
  # scale_fill_manual("", values = c("fit" = "black", "prior-predict" = "darkgrey")) +
  # facet_wrap(vars(fit_type), ncol = 1) +
  coord_cartesian(xlim = c(0, 2)) +
  NULL
```

```{r, eval=FALSE}
dist_fit_data %>%
  # filter(fct_match(fit_type, "prior-predict")) %>% 
  select(model, model_name, mu_rep, fit_type) %>% 
  unnest(mu_rep) %>%
  plot_estimand_hist(iter_est, binwidth = 0.02, results_group = fit_type, group_labels = c(fit = "Posterior", "prior-predict" = "Prior")) +
  scale_x_continuous() +
  scale_y_continuous("") +
  labs(
    title = latex2exp::TeX(r"{Signaling $\mu$}"),
    x = ""
  ) +
  facet_wrap(vars(assigned_treatment), labeller = as_labeller(str_to_title)) +
  coord_cartesian(xlim = c(0, 5)) +
  theme(legend.position = "none") +
  NULL
```

```{r, eval=FALSE}
dist_prior <- dir("data/stan_analysis_data/", pattern = str_glue(r"{dist_prior{params$structural_takeup_version}_STRUCTURAL_LINEAR_U_SHOCKS-\d\.csv}"), full.names = TRUE) %>% 
  as_cmdstan_fit()

dist_fit <- dir("data/stan_analysis_data/", pattern = str_glue(r"{dist_fit{params$structural_takeup_version}_STRUCTURAL_LINEAR_U_SHOCKS-\d\.csv}"), full.names = TRUE) %>% 
  as_cmdstan_fit()

all_fit <- lst(
  prior = dist_prior,
  fit = dist_fit
) 
```

```{r, eval=FALSE}
all_fit %>% 
  map(posterior::as_draws_df) %>% 
  map_dfr(spread_draws, raw_u_sd[treatment], .id = "fit_type") %>% 
  ggplot() +
  geom_halfeyeh(aes(raw_u_sd, fill = fit_type), alpha = 0.25)
```


```{r, fig.height=1.8, fig.width=8, eval=FALSE}
plot_pos <- position_dodge(width = 0.75)

all_fit %>% 
  map_dfr(~ .x$draws(c("beta_control", "beta_ink_effect", "beta_bracelet_effect", "wtp_value_utility")) %>% 
            posterior::as_draws_df() %>% 
            mutate(iter_id = .draw) %>% 
            pivot_longer(!c(iter_id, .draw, .chain, .iteration), names_to = "param", values_to = "iter_value") %>% 
            nest(iter_data = c(.chain, .iteration, .draw, iter_id, iter_value)), 
          .id = "fit_type") %>% 
  mutate(quant = map(iter_data, quantilize_est, iter_value, quant_probs = c(0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95))) %>% 
  unnest(quant) %>% 
  ggplot(aes(y = param, group = fit_type)) +
  geom_linerange(aes(xmin = per_0.25, xmax = per_0.75, color = fit_type), alpha = 0.4, size = 3, position = plot_pos) +
  geom_crossbar(aes(x = per_0.5, xmin = per_0.1, xmax = per_0.9, color = fit_type), fatten = 0, size = 0.4, width = 0.5, position = plot_pos) +
  geom_linerange(aes(xmin = per_0.05, xmax = per_0.95, color = fit_type), size = 0.4, position = plot_pos) +
  NULL
```

### Posterior Predictive Checking

```{r, fig.width=7.5, fig.height=8}
#| pp-check
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want)) %>% 
  select(model_name, cluster_cf_cutoff) %>% 
  unnest(cluster_cf_cutoff) %>% 
  filter(!is.na(obs_prop_takeup), fct_match(assigned_treatment, c("control", "bracelet"))) %>% 
  ggplot(aes(assigned_dist_obs)) +
  geom_point(aes(y = obs_prop_takeup, color = assigned_treatment), size = 0.8, show.legend = FALSE) + #, data = . %>% filter(fct_match(model_type, "structural"))) +
  # geom_pointrange(aes(y = per_0.5, ymin = per_0.1, ymax = per_0.9, color = assigned_treatment), # color = model_type, group = model_type), 
  #                 fatten = 0.8, position = position_dodge(width = 50)) +
  geom_crossbar(aes(y = per_0.5, ymin = per_0.1, ymax = per_0.9, color = assigned_treatment), width = 40) +
  labs(x = "Distance", y = "Take-up") +
  # facet_wrap(vars(assigned_treatment), ncol = 1) +
  facet_wrap(vars(model_name), ncol = 1) +
  theme(legend.position = "bottom")
```

```{r, fig.width=8, fig.height=8}
#| other-pp-check
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, models_we_want)) %>% 
  select(model_type, model_name, cluster_cf_cutoff) %>% 
  unnest(cluster_cf_cutoff) %>% 
  filter(!is.na(obs_prop_takeup), fct_match(assigned_treatment, c("control", "ink", "bracelet"))) %>% 
  ggplot(aes(assigned_dist_obs)) +
  geom_point(aes(y = obs_prop_takeup), size = 0.8, show.legend = FALSE) + 
  geom_ribbon(aes(ymax = per_0.9, ymin = per_0.1), alpha = 0.25) +
  geom_line(aes(y = per_0.5), alpha = 0.25) +
  labs(
    title = "Posterior Checking",
    x = "Distance to Treatment", y = "Proportion of Takeup", 
    caption = "Dots: observed proportion of takeup. Line and ribbon: estimated proportion of takeup (median and 80% CI)."
  ) +
  facet_grid(rows = vars(model_name), cols = vars(assigned_treatment), labeller = labeller(.default = str_to_title)) +
  theme(legend.position = "bottom")
```

# Willingness-to-Pay

```{r, eval=FALSE}
wtp_results <- read_rds(file.path("data", "stan_analysis_data", "wtp_model_fit_results.rds")) 
wtp_prior_results <- read_rds(file.path("data", "stan_analysis_data", "wtp_model_prior_results.rds"))
```

```{r wtp-post}
joint_wtp_results <- dist_fit_data %>%
  filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS_NO_BELIEFS_DIST")) %>%
  pull(wtp_results) %>%
  first()
if (!is.null(joint_wtp_results)) {

joint_wtp_results %>% 
  mutate(fit_type = "joint") %>% 
  plot_wtp_results()
}
```

```{r wtp-prior}
dist_fit_data %>% 
  filter(fct_match(fit_type, "prior-predict"), fct_match(model_type, "structural")) %>% 
  pull(wtp_results) %>%
  first() %>%
  mutate(fit_type = "joint") %>% 
  plot_wtp_results()
```

```{r wtp-post-compare, eval=FALSE}
bind_rows(
  joint = joint_wtp_results,
  alone = wtp_results, 
  .id = "fit_type"
) %>%
  plot_wtp_results()
```

# Beliefs

```{r, eval=FALSE}
beliefs_results <- read_rds(file.path("data", "stan_analysis_data", "beliefs_results.rds"))
beliefs_prior_results <- read_rds(file.path("data", "stan_analysis_data", "beliefs_prior_results.rds"))
```

```{r beliefs-line-plot-post-alone, fig.width=8, fig.height=8, eval=FALSE}
beliefs_results %>%
  plot_beliefs_est("Posterior Estimation")
```


```{r beliefs-line-plot-post, fig.width=8, fig.height=8}
dist_fit_data %>%
  filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP")) %>%
  pull(beliefs_results) %>%
  first() %>%
# beliefs_results %>%
  plot_beliefs_est("Posterior Estimation")

# dist_fit_data %>%
#   filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS_NO_BELIEFS_DIST")) %>%
#   pull(beliefs_results) %>%
#   first() %>%
# # beliefs_results %>%
#   plot_beliefs_est("Posterior Estimation")
```

```{r beliefs-line-plot-prior, fig.width=8, fig.height=8}
dist_fit_data %>%
  filter(fct_match(fit_type, "prior-predict"), fct_match(model_type, "structural")) %>%
  pull(beliefs_results) %>%
  first() %>%
# beliefs_results %>% 
  plot_beliefs_est("Prior Prediction")
```

```{r beliefs-hist-post, fig.width=8, fig.height=8}
plot_beliefs_hist <- function(beliefs_results, top_title, binwidth = NULL) {
  pos_dodge <- position_dodge(width = 0.3)

  first_plot <- beliefs_results$prob_knows %>% 
    filter(ord == 1) %>% 
    ggplot(aes(group = assigned_dist_group)) +
    geom_histogram(aes(x = iter_est, y = stat(density) * (binwidth %||% 1), fill = assigned_dist_group),
                   binwidth = binwidth, position = "identity", alpha = 0.5,
                   data = . %>% unnest(iter_data)) + 
    scale_fill_canva("", labels = str_to_title, palette = canva_palette_vibrant) + 
    labs(
      title = "First Order Beliefs",
      subtitle = "Proportion",
      x = "", y = "") +
    facet_wrap(vars(assigned_treatment), ncol = 1) +
    theme(legend.position = "top", axis.text.y = element_blank()) +
    NULL
  
  cowplot::plot_grid(
    cowplot::ggdraw() +
      cowplot::draw_label(top_title, size = 20, fontface = "italic"),
    cowplot::plot_grid(
      first_plot +
        theme(
          legend.position = "none"
        ) +
        NULL,
      
      beliefs_results$ate_knows %>%
        filter(ord == 1, assigned_dist_group_left == assigned_dist_group_right) %>%
        ggplot(aes(group = assigned_dist_group_left)) +
        geom_histogram(aes(x = iter_est, y = stat(density) * (binwidth %||% 1), fill = assigned_dist_group_left),
                       binwidth = binwidth, position = "identity", alpha = 0.5,
                       data = . %>% unnest(iter_data)) + 
        scale_fill_canva("", labels = str_to_title, palette = canva_palette_vibrant) + 
        labs(
          title = "",
          subtitle = "Treatment Effect",
          x = "", y = "") +
        facet_wrap(vars(assigned_treatment_left), ncol = 1) +
        theme(legend.position = "none", axis.text.y = element_blank()) +
        NULL,
      
      beliefs_results$prob_knows %>% 
        filter(ord == 2) %>% 
        ggplot(aes(group = assigned_dist_group)) +
        geom_histogram(aes(x = iter_est, y = stat(density) * (binwidth %||% 1), fill = assigned_dist_group),
                       binwidth = binwidth, position = "identity", alpha = 0.5,
                       data = . %>% unnest(iter_data)) + 
        scale_fill_canva("", labels = str_to_title, palette = canva_palette_vibrant) + 
        labs(
          title = "Second Order Beliefs",
          subtitle = "Proportion",
          x = "", y = "") +
        facet_wrap(vars(assigned_treatment), ncol = 1) +
        theme(legend.position = "none", axis.text.y = element_blank()) +
        NULL,
      
      beliefs_results$ate_knows %>%
        filter(ord == 2, assigned_dist_group_left == assigned_dist_group_right) %>%
        ggplot(aes(group = assigned_dist_group_left)) +
        geom_histogram(aes(x = iter_est, y = stat(density) * (binwidth %||% 1), fill = assigned_dist_group_left),
                       binwidth = binwidth, position = "identity", alpha = 0.5,
                       data = . %>% unnest(iter_data)) + 
        scale_fill_canva("", labels = str_to_title, palette = canva_palette_vibrant) + 
        labs(
          title = "",
          subtitle = "Treatment Effect",
          x = "", y = "") +
        facet_wrap(vars(assigned_treatment_left), ncol = 1) +
        theme(legend.position = "none", axis.text.y = element_blank()) +
        NULL,
      
      ncol = 2, axis = "b", align = "h" 
    ),
    
    cowplot::get_legend(first_plot),
    
    ncol = 1, rel_heights = c(0.1, 1, 0.1)
  )
}

dist_fit_data %>%
  filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS_PHAT_MU_REP")) %>%
  pull(beliefs_results) %>%
  first() %>%
# beliefs_prior_results %>% 
  plot_beliefs_hist("Posterior Estimation", binwidth = 0.01)

# dist_fit_data %>%
#   filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS_NO_BELIEFS_DIST")) %>%
#   pull(beliefs_results) %>%
#   first() %>%
# # beliefs_prior_results %>% 
#   plot_beliefs_hist("Posterior Estimation", binwidth = 0.01)
```

```{r beliefs-hist-prior, fig.width=8, fig.height=8}
dist_fit_data %>%
  filter(fct_match(fit_type, "prior-predict"), fct_match(model_type, "structural")) %>%
  pull(beliefs_results) %>%
  first() %>%
# beliefs_prior_results %>% 
  plot_beliefs_hist("Prior Prediction", binwidth = 0.01)
```

# Distance Model

```{r, eval=FALSE}
alone_dist_fit_data <- dist_fit_data

alone_dist_fit_data$group_dist_param[[1]] <- read_rds("data/stan_analysis_data/dist_model_fit_fit.rds") %>% 
  get_dist_results(dist_fit_data$stan_data[[1]])

alone_dist_fit_data$group_dist_param[[2]] <- read_rds("data/stan_analysis_data/dist_model_prior_prior.rds") %>% 
  get_dist_results(dist_fit_data$stan_data[[1]])
```

```{r}
dist_fit_data$stan_data %>% 
  first() %>% 
  pluck("analysis_data") %>% 
  distinct(county, cluster_id, assigned_dist_group = dist.pot.group, cluster.dist.to.pot) %>% 
  ggplot() +
  # geom_freqpoly(aes(x = cluster.dist.to.pot, y = stat(density)), binwidth = 300) +
  geom_histogram(aes(x = cluster.dist.to.pot, y = stat(density) * 300), alpha = 0.5, binwidth = 400, boundary = 0) +
  labs(x = "", y = "") +
  facet_grid(rows = vars(county), cols = vars(assigned_dist_group), labeller = labeller(.cols = str_to_title)) +
  NULL
```

```{r}
plot_dist_model <- function(results, plot_obs_dist = TRUE, subsample_size = 1500, x_range = seq(-500, 3000, 10)) {
  plot_obj <- results %>% 
    select(stan_data, group_dist_param) %>% 
    mutate(distrib_samples = map2(stan_data, group_dist_param, ~ {
      subsample_ids <- if (is_null(subsample_size)) seq(max(.y$iter_id)) else sample(max(.y$iter_id), subsample_size)
     
      filter(.y, iter_id %in% subsample_ids) %>%  
        mutate(
          samples = map2(group_dist_mean, group_dist_sd, ~ {
            tibble(x = x_range,
                   # density = dnorm(x, mean = ..1 * ..3, sd = ..2 * ..3))
                   density = dnorm(log(x / ..3), mean = ..1, sd = ..2) / x)
          },
          sd(.x$analysis_data$cluster.dist.to.pot))) %>% 
        unnest(samples) %>% 
        group_by(assigned_dist_group, iter_id, x) %>% 
        summarize(
          density = weighted.mean(density, group_dist_mix),
          .groups = "drop"
        )
    })) %>% 
    ggplot() +
    geom_line(aes(x, density, group = iter_id, color = "Posterior Samples"), alpha = 0.04,
              data = . %>% 
                select(distrib_samples) %>% 
                unnest(distrib_samples)
              ) +
    scale_color_manual("", values = c("Posterior Samples" = "black", "Observed Sample Frequency" = "firebrick3")) +
    labs(
      title = "Fit and observed distance to deworming treatment, conditional on assigned distance group (g).",
      x = "Distance",
      y = "",
      caption = "The observed sample distance frequency plot was calculated using bins of width 200 meters.
                 Posterior sample distributions are shown using 500 iterations."
    ) +
    coord_cartesian(xlim = c(0, 3000), ylim = c(0, 0.0025)) +
    # theme(axis.text.y = element_blank()) +
    facet_wrap(vars(assigned_dist_group), ncol = 1, labeller = as_labeller(. %>% str_c("g = ", .))) +
    NULL
  
  if (plot_obs_dist) {
    plot_obj <- plot_obj +
      geom_freqpoly(aes(x = cluster.dist.to.pot, y = stat(density), color = "Observed Sample Frequency"),
                     binwidth = 200,
                     size = 0.8,
                     data = . %>% 
                      pluck("stan_data") %>% 
                      first() %>% 
                      pluck("analysis_data") %>% 
                      distinct(cluster_id, assigned_dist_group = dist.pot.group, cluster.dist.to.pot)) 
  }
  
  return(plot_obj)
}
```

```{r dist-post-alone, fig.width=10, eval=FALSE}
alone_dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model_type, "structural")) %>% 
  plot_dist_model(subsample_size = NULL)
```

```{r dist-post, fig.width=10, cache=TRUE, eval=FALSE}
dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS")) %>%
  plot_dist_model(subsample_size = NULL)

dist_fit_data %>% 
  filter(fct_match(fit_type, "fit"), fct_match(model, "STRUCTURAL_LINEAR_U_SHOCKS_NO_BELIEFS_DIST")) %>%
  plot_dist_model(subsample_size = NULL)
```

```{r dist-prior, fig.width=10, cache=TRUE}, eval=FALSE
dist_fit_data %>% 
  filter(fct_match(fit_type, "prior-predict"), fct_match(model_type, "structural")) %>% 
  plot_dist_model(plot_obs_dist = FALSE, subsample_size = NULL) + coord_cartesian(xlim = NULL, ylim = c(0, 0.0025))
```

```{r, fig.width=10, eval=FALSE}
read_rds("data/stan_analysis_data/dist_model_fit_fit.rds") %>%
  get_imputed_dist(dist_fit_data$stan_data[[1]]) %>%
  filter(!is.na(obs_standard_dist)) %>% 
  mutate(cluster_id = factor(cluster_id) %>% fct_reorder(obs_standard_dist)) %>% 
  ggplot(aes(cluster_id)) +
  geom_point(aes(y = obs_standard_dist)) +
  geom_linerange(aes(ymin = per_0.1, ymax = per_0.9)) +
  facet_grid(rows = vars(dist_treatment), cols = vars(county), scales = "free_x")
```

# Reduced Form

```{r reduced-form-incentive-te, fig.width=8, fig.height=9}
dist_fit_data %>%
  filter(fct_match(model_type, "reduced form")) %>% 
  mutate(est_takeup_te =
    map(est_takeup_te, 
        filter, 
        (is.na(assigned_dist_group_left) & is.na(assigned_dist_group_right)) | (assigned_dist_group_left == assigned_dist_group_right),
        assigned_treatment_left != assigned_treatment_right,
        fct_match(assigned_treatment_right, c("control")),
        fct_match(assigned_treatment_left, "bracelet") | !fct_match(assigned_treatment_right, "calendar")) %>% 
      map(
        mutate,
        assigned_dist_group_left = fct_explicit_na(assigned_dist_group_left, "Combined") %>% 
          fct_relabel(str_to_title) %>% 
          fct_relevel("Combined"),
        assigned_treatment_left = str_to_title(assigned_treatment_left)
      )
  ) %>% 
  select(model, model_name, est_takeup_te, fit_type, model_color) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_te, assigned_treatment_left) +
        scale_y_discrete("") +
        labs(
          title = "Incentive Average Treatment Effect"
        ) +
        ggforce::facet_col(vars(assigned_dist_group_left), space = "free", scales = "free_y") +
        theme(legend.position = "none") +
        NULL,
      
      plot_estimand_hist(., est_takeup_te, iter_takeup_te, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.2)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group_left), rows = vars(assigned_treatment_left)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1
    )
  } 
```

```{r, fig.width=8, fig.height=11}
dist_fit_data %>% 
  filter(fct_match(model_type, "reduced form")) %>% 
  select(model, model_name, est_takeup_level, fit_type, model_color) %>% 
  mutate(
    est_takeup_level = map(
      est_takeup_level,
      mutate,
      assigned_dist_group = fct_explicit_na(assigned_dist_group, "Combined") %>% 
        fct_relabel(str_to_title) %>% 
        fct_relevel("Combined"),
      assigned_treatment = str_to_title(assigned_treatment),
    )
  ) %>% {
    cowplot::plot_grid(
      plot_estimands(., est_takeup_level, assigned_treatment) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_discrete("") +
        labs(
          title = "Incentive Take-up Level"
        ) +
        ggforce::facet_col(vars(assigned_dist_group), space = "free", scales = "free_y") +
        theme(legend.position = "none") +
        NULL,
      
      plot_estimand_hist(., est_takeup_level, iter_prop_takeup, binwidth = 0.05) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_continuous("") +
        facet_grid(cols = vars(assigned_dist_group), rows = vars(assigned_treatment), scales = "free_y") +
        theme(legend.position = "none") +
        coord_cartesian(ylim = c(0, 0.2)) +
        NULL,
      
      ncol = 1
    )
  }
```

```{r reduced-form-dist-te, fig.width=8, fig.height=8}
dist_fit_data %>%
  filter(fct_match(model_type, "reduced form")) %>% 
  mutate(est_takeup_te =
    map(est_takeup_te, 
        filter, 
        !is.na(assigned_dist_group_left) & !is.na(assigned_dist_group_right),
        fct_match(assigned_dist_group_left, "far"), 
        fct_match(assigned_dist_group_right, "close"), 
        assigned_treatment_left == assigned_treatment_right) %>% 
      map(
        mutate,
        assigned_treatment_left = fct_relabel(assigned_treatment_left, str_to_title)
      )
  ) %>% 
  select(model, model_name, est_takeup_te, fit_type, model_color) %>% { 
    cowplot::plot_grid(
      plot_estimands(., est_takeup_te, assigned_treatment_left) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.1)) +
        scale_y_discrete("") +
        labs(
          title = "Distance Average Treatment Effect"
        ) +
        theme(legend.position = "none") +
        NULL,
      
      plot_estimand_hist(., est_takeup_te, iter_takeup_te, binwidth = 0.02) +
        scale_x_continuous("", breaks = seq(-1, 1, 0.2)) +
        scale_y_continuous("") +
        facet_wrap(vars(assigned_treatment_left)) +
        theme(legend.position = "none") +
        NULL,
      
      ncol = 1, rel_heights = c(0.6, 1)
    )
  } 
```

## By Age Group

```{r}
rf_analysis_data %>% 
  ggplot(aes(age.census_group)) +
  geom_bar(alpha = 0.5) +
  labs(title = "Age Group Sizes", x = "Age Group", y = "Size")
```

```{r, fig.width=10, eval=FALSE}
obs_takeup_by_age <- dist_fit_data %>% 
  filter(
    fct_match(model_type, "reduced form"),
    fct_match(fit_type, "fit"),
  ) %>% 
  select(obs_age_group_cluster_takeup_level) %>% 
  unnest(obs_age_group_cluster_takeup_level) %>%  
  group_by(age_group, assigned_treatment_obs, assigned_dist_group_obs) %>%
  summarize(obs_takeup_prop = sum(obs_num_takeup) / sum(cluster_size), .groups = "drop")
# rf_analysis_data %>%
#   group_by(age_group = age.census_group, assigned_treatment, assigned_dist_group) %>% 
#   summarize(obs_takeup_prop = mean(dewormed), .groups = "drop") %>% 

obs_takeup_by_age %>% 
  ggplot(aes(age_group, obs_takeup_prop)) +
  geom_col(aes(fill = assigned_treatment_obs), position = "dodge", alpha = 0.5) +
  scale_fill_discrete("", labels = str_to_title) +
  labs(title = "Takeup by Age Group", x = "Age Group", y = "Size") +
  facet_wrap(vars(assigned_dist_group_obs), labeller = labeller(.default = str_to_title))
```


```{r, fig.width=10, fig.height=8, eval=FALSE}
age_group_te_iter_data <- dist_fit_data %>% 
  filter(
    fct_match(model_type, "reduced form"),
    # fct_match(fit_type, "prior-predict"),
  ) %>% 
  select(model, model_name, est_age_group_takeup_te, fit_type) %>% 
  unnest(est_age_group_takeup_te) %>% 
  filter(
    (is.na(assigned_dist_group_left) & is.na(assigned_dist_group_right)) | (assigned_dist_group_left == assigned_dist_group_right),
    assigned_treatment_left != assigned_treatment_right,
    fct_match(assigned_treatment_right, c("control")),
    fct_match(assigned_treatment_left, "bracelet") | !fct_match(assigned_treatment_right, "calendar")
  ) %>% 
  mutate(
    assigned_dist_group_left = fct_explicit_na(assigned_dist_group_left, "Combined") %>% 
      fct_relabel(str_to_title) %>% 
      fct_relevel("Combined"),
    assigned_treatment_left = str_to_title(assigned_treatment_left)
  ) %>%
  select(!c(mean_est, starts_with("per_"))) %>% 
  unnest(iter_data)
  
age_group_te_iter_data %>% 
  ggplot(aes(age_group, iter_takeup_te)) +
  # geom_line(aes(color = fit_type, group = str_c(fit_type, iter_id)), alpha = 0.1, data = . %>% filter(iter_id <= 100)) +
  stat_lineribbon(aes(fill = fit_type, color = fit_type), size = 0.5, .width = c(0.5, 0.8), alpha = 0.25) +
  geom_hline(yintercept = 0, linetype = "dotted") +
  scale_fill_discrete("", labels = c("fit" = "Posterior", "prior-predict" = "Prior Predict"), aesthetics = c("fill", "color")) +
  labs(title = "Takeup Treatment Effect by Age Group", x = "Age Group", y = "") +
  facet_grid(vars(assigned_dist_group_left), vars(assigned_treatment_left)) +
  theme(legend.position = "top")
```

```{r, fig.width=10, fig.height=8, eval=FALSE}
dist_fit_data %>% 
  filter(
    fct_match(model_type, "reduced form"),
    # fct_match(fit_type, "prior-predict")
  ) %>% 
  select(model, model_name, est_age_group_takeup_level, fit_type) %>% 
  unnest(est_age_group_takeup_level) %>% 
  mutate(
    assigned_dist_group = fct_explicit_na(assigned_dist_group, "combined") %>% 
      # fct_relabel(str_to_title) %>% 
      fct_relevel("combined"),
  ) %>% 
  select(!c(mean_est, starts_with("per_"))) %>% 
  unnest(iter_data) %>% 
  ggplot(aes(age_group, iter_prop_takeup)) +
  # geom_line(aes(color = fit_type, group = str_c(fit_type, iter_id)), alpha = 0.1, data = . %>% filter(iter_id <= 100)) +
  stat_lineribbon(aes(fill = fit_type, color = fit_type), size = 0.5, .width = c(0.5, 0.8), alpha = 0.25) +
  geom_point(aes(y = obs_takeup_prop), size = 0.5, data = obs_takeup_by_age %>% rename_with(str_remove, ends_with("_obs"), "_obs$")) +
  scale_fill_discrete("", labels = c("fit" = "Posterior", "prior-predict" = "Prior Predict"), aesthetics = c("fill", "color")) +
  labs(title = "Takeup Levels by Age Group", x = "Age Group", y = "", subtitle = "Black points: observed levels.") +
  facet_grid(vars(assigned_dist_group), vars(assigned_treatment), labeller = labeller(.default = str_to_title)) +
  theme(legend.position = "top")
```

## Over Time

```{r, fig.width=8}
rf_analysis_data %>% 
  count(assigned_treatment, dewormed, dewormed.day) %>% 
  group_by(assigned_treatment) %>% 
  mutate(p = n / sum(n)) %>% 
  ungroup() %>% 
  filter(
    !is.na(dewormed.day),
    # fct_match(assigned_treatment, c("control", "bracelet"))
  ) %>% 
  ggplot(aes(dewormed.day, p)) +
  geom_step(aes(color = assigned_treatment), data = . %>% filter(!fct_match(assigned_treatment, "control"))) +
  geom_step(linetype = "dotted", size = 1, alpha = 0.5, data = . %>% filter(fct_match(assigned_treatment, "control"))) +
  scale_x_continuous("Deworming Day", breaks = 1:12) +
  scale_color_discrete("", labels = str_to_title) +
  labs(title = "Proportion Dewormed by Day", y = "", caption = "Dotted line is the control proportion.")
```