library(dplyr)
library(readr)
library(purrr)
library(tidyr)

library(brms)
library(ggdist)

library(ggplot2)
theme_set(theme_classic() + theme(strip.background = element_rect(fill = "white"), plot.background = element_blank()))
library(cowplot)

cong_cols <- c("#E69F00", "#009E73")

norm01 <- function(x, ...) (x-min(x, ...)) / (max(x, ...) - min(x, ...))
norm01_manual <- function(x, min_x, max_x) (x-min_x) / (max_x - min_x)

# summarise posteriors from behavioural validation experiment -------------

# valid_m <- readRDS(file.path("..", "01 Validation", "02 Analysis", "mods", "m_bme.rds"))
# 
# valid_m_ests <- valid_m %>%
#   as_draws_df("^b\\_|^sd\\_|^cor\\_", regex=TRUE) %>%
#   select(-starts_with(".")) %>%
#   pivot_longer(cols=everything(), names_to="par", values_to="est")
# 
# lapply(unique(valid_m_ests$par), function(p) {
#   d_i <- filter(valid_m_ests, par == p)
#   
#   ggplot(d_i, aes(est)) +
#     geom_function(fun=dnorm, args=list(mean=mean(d_i$est), sd=sd(d_i$est)), colour="red") +
#     geom_density() +
#     geom_vline(xintercept = median(d_i$est), colour="blue") +
#     facet_wrap(vars(par), scales="free") +
#     labs(x=NULL, y=NULL)
# }) %>%
#   wrap_plots()
# 
# valid_m_norm <- valid_m_ests %>%
#   group_by(par) %>%
#   summarise(
#     median_est = median(est),
#     mean_est = mean(est),
#     sd_est = sd(est),
#     sd_est10 = sd_est*10
#   )
# 
# rm(valid_m)
# gc()

# import data -------------------------------------------------------------

# get the stimuli's percentage of name agreement values
stim <- read_csv("boss.csv", col_types = cols(perc_name_agree_denom_fq_inputs = col_number())) %>%
  select(filename, perc_name_agree_denom_fq_inputs) %>%
  rename(perc_name_agree = perc_name_agree_denom_fq_inputs)

d <- file.path("raw_data", "stim-pc", "data", "pictureword") %>%
  list.files(pattern = "^.*\\.csv$", full.names = TRUE) %>%
  map_df(read_csv, col_types = cols(sex="c")) %>%
  filter(acc == 1) %>%
  left_join(stim, by=c("image" = "filename")) %>%
  mutate(
    prop_agree = perc_name_agree/100,
    pred_norm = norm01(prop_agree),
    cong_dev = scale(if_else(condition == "A1", 1, 0), center = TRUE, scale = FALSE)
  )

# setup priors for RT model -----------------------------------------------

priors <- c(
  # FIXED EFFECTS
  #  mu
  set_prior("normal(5.75, 0.71)", class = "b", coef = "Intercept"),
  set_prior("normal(0.472, 0.875)", class = "b", coef = "cong_dev"),
  set_prior("normal(-0.543, 0.78)", class = "b", coef = "pred_norm"),
  set_prior("normal(-0.671, 1.29)", class = "b", coef = "cong_dev:pred_norm"),
  #  sigma
  set_prior("normal(-0.85, 0.535)", class = "b", coef = "Intercept", dpar="sigma"),
  set_prior("normal(0.0404, 0.94)", class = "b", coef = "cong_dev", dpar="sigma"),
  set_prior("normal(0.229, 0.755)", class = "b", coef = "pred_norm", dpar="sigma"),
  set_prior("normal(0.142, 1.345)", class = "b", coef = "cong_dev:pred_norm", dpar="sigma"),
  #  delta
  set_prior("normal(0, 7.5)", class = "b", coef = "Intercept", dpar="ndt"),  # wider than other priors, and equivalent to a delay of just exp(3) = 20 ms rather than the exp(5.63) from the validation posterior, because I expected the forced delay of response (until colour change) to greatly reduce non-decision time, but I'm not sure by how much
  set_prior("normal(-0.4, 7.5)", class = "b", coef = "cong_dev", dpar="ndt"),
  set_prior("normal(0.132, 7.5)", class = "b", coef = "pred_norm", dpar="ndt"),
  set_prior("normal(-0.671, 7.5)", class = "b", coef = "cong_dev:pred_norm", dpar="ndt"),
  # STANDARD DEVIATIONS OF RANDOM EFFECT DISTRIBUTIONS
  #  mu
  #    -subj_id
  set_prior("student_t(10, 0.29, 0.1)", class = "sd", coef = "Intercept", group = "subj_id"),
  set_prior("student_t(10, 0.079, 0.1)", class = "sd", coef = "cong_dev", group = "subj_id"),
  set_prior("student_t(10, 0.128, 0.1)", class = "sd", coef = "pred_norm", group = "subj_id"),
  set_prior("student_t(10, 0.077, 0.25)", class = "sd", coef = "cong_dev:pred_norm", group = "subj_id"),
  #    -image
  set_prior("student_t(10, 0.116, 0.05)", class = "sd", coef = "Intercept", group = "image"),
  set_prior("student_t(10, 0.137, 0.1)", class = "sd", coef = "cong_dev", group = "image"),
  #    -string
  set_prior("student_t(10, 0.379, 0.05)", class = "sd", coef = "Intercept", group = "string"),
  #  sigma
  #    -subj_id
  set_prior("student_t(10, 0.98, 0.1)", class = "sd", coef = "Intercept", group = "subj_id", dpar = "sigma"),
  set_prior("student_t(10, 0.121, 0.1)", class = "sd", coef = "cong_dev", group = "subj_id", dpar = "sigma"),
  set_prior("student_t(10, 0.075, 0.1)", class = "sd", coef = "pred_norm", group = "subj_id", dpar = "sigma"),
  set_prior("student_t(10, 0.084, 0.25)", class = "sd", coef = "cong_dev:pred_norm", group = "subj_id", dpar = "sigma"),
  #    -image
  set_prior("student_t(10, 0.068, 0.05)", class = "sd", coef = "Intercept", group = "image", dpar = "sigma"),
  set_prior("student_t(10, 0.1, 0.1)", class = "sd", coef = "cong_dev", group = "image", dpar = "sigma"),
  #    -string
  set_prior("student_t(10, 0.039, 0.05)", class = "sd", coef = "Intercept", group = "string", dpar = "sigma"),
  #  delta
  set_prior("student_t(10, 0.096, 0.1)", class = "sd", coef = "Intercept", group = "subj_id", dpar = "ndt"),
  set_prior("student_t(10, 0.071, 0.1)", class = "sd", coef = "cong_dev", group = "subj_id", dpar = "ndt"),
  set_prior("student_t(10, 0.028, 0.1)", class = "sd", coef = "pred_norm", group = "subj_id", dpar = "ndt"),
  set_prior("student_t(10, 0.038, 0.25)", class = "sd", coef = "cong_dev:pred_norm", group = "subj_id", dpar = "ndt"),
  #    -image
  set_prior("student_t(10, 0.245, 0.05)", class = "sd", coef = "Intercept", group = "image", dpar = "ndt"),
  set_prior("student_t(10, 0.023, 0.1)", class = "sd", coef = "cong_dev", group = "image", dpar = "ndt"),
  #    -string
  set_prior("student_t(10, 0.015, 0.05)", class = "sd", coef = "Intercept", group = "string", dpar = "ndt")
)

n_cores <- 7

seed <- 3101
n_iter <- 10000
n_warmup <- 7500
adapt_delta <- 0.99
max_treedepth <- 10
n_chains <- 5
refresh <- 100

f <- brmsformula(
  rt ~ 0 + Intercept + cong_dev * pred_norm +
    (cong_dev * pred_norm | subj_id) +
    (cong_dev | image) +
    (1 | string),
  sigma ~ 0 + Intercept + cong_dev * pred_norm +
    (cong_dev * pred_norm | subj_id) +
    (cong_dev | image) +
    (1 | string),
  ndt ~ 0 + Intercept + cong_dev * pred_norm +
    (cong_dev * pred_norm | subj_id) +
    (cong_dev | image) +
    (1 | string)
)

m_rt <- brm(
  formula = f,
  data = d,
  family = shifted_lognormal(),
  prior = priors,
  iter = n_iter,
  warmup = n_warmup,
  chains = n_chains,
  control = list(
    adapt_delta = adapt_delta,
    max_treedepth = max_treedepth
  ),
  init = replicate(
    n_chains,
    list(b_ndt = as.array(rep(-5, 4))),
    simplify=FALSE
  ),
  sample_prior = "no",
  silent = TRUE,
  cores = n_cores,
  seed = seed,
  thin = 1,
  file = file.path("mods", "m_rt.rds"),
  refresh = refresh
)

# plot results ------------------------------------------------------------

# get predicted densities
coding_lookup <- d %>%
  group_by(condition) %>%
  summarise(cong_dev = unique(cong_dev))

props <- 1:10/10

fe_tidy <- fixef(m_rt, robust=TRUE) %>%
  as_tibble(rownames="term")

fe <- sapply(fe_tidy$term, function(term_i) {
  fe_tidy %>%
    filter(term==term_i) %>%
    pull(Estimate)
})

fe_preds <- tibble(
  condition = rep(c("A1", "A2"), each = length(props)),
  condition_label = if_else(condition=="A1", "Congruent", "Incongruent"),
  prop_agree = rep(props, 2)
) %>%
  left_join(coding_lookup, by = "condition") %>%
  mutate(
    pred_norm = norm01_manual(prop_agree, min(d$prop_agree), max(d$prop_agree)),
    int_mu = fe["Intercept"],
    int_sigma = fe["sigma_Intercept"],
    int_ndt = fe["ndt_Intercept"],
    cong_mu = fe["cong_dev"],
    cong_sigma = fe["sigma_cong_dev"],
    cong_ndt = fe["ndt_cong_dev"],
    pred_norm_mu = fe["pred_norm"],
    pred_norm_sigma = fe["sigma_pred_norm"],
    pred_norm_ndt = fe["ndt_pred_norm"],
    interact_mu = fe["cong_dev:pred_norm"],
    interact_sigma = fe["sigma_cong_dev:pred_norm"],
    interact_ndt = fe["ndt_cong_dev:pred_norm"],
    pred_mu = int_mu + cong_dev*cong_mu + pred_norm*pred_norm_mu + cong_dev*pred_norm*interact_mu,
    pred_sigma = int_sigma + cong_dev*cong_sigma + pred_norm*pred_norm_sigma + cong_dev*pred_norm*interact_sigma,
    pred_ndt = int_ndt + cong_dev*cong_ndt + pred_norm*pred_norm_ndt + cong_dev*pred_norm*interact_ndt
  )

quantities <- 0:1000

fe_cond_dens <- map_dfr(quantities, function(q) mutate(fe_preds, rt = q)) %>%
  mutate(
    pred_dens = dshifted_lnorm(
      x = rt,
      meanlog = pred_mu,
      sdlog = exp(pred_sigma),
      shift = exp(pred_ndt)
    )
  )


# build panel A

panel_A_margin <- theme_get()$plot.margin
panel_A_margin[[2]] <- unit(0.2, "npc")

pub_panel_A <- fe_cond_dens %>%
  mutate(condition_label = sprintf("Picture-%s", condition_label)) %>%
  ggplot(aes(rt, pred_dens, colour = prop_agree)) +
  geom_line(aes(group = as.factor(prop_agree))) +
  facet_wrap(~condition_label) +
  labs(x = "Response Time (ms)", y = "Predicted Density", colour = "Predictability", tag="a") +
  scale_colour_continuous(
    type="viridis", breaks=sort(unique(fe_cond_dens$prop_agree)),
    labels=sprintf("%s%%", sort(unique(fe_cond_dens$prop_agree))*100),
    # guide=guide_colourbar(barheight = 7.175)
    guide = guide_legend(override.aes = list(linewidth=1), reverse = TRUE)
  ) +
  theme_classic() +
  theme(
    plot.margin = panel_A_margin,
    legend.position = c(1.15, 0.5875),
    legend.key.height = unit(11, "pt"),
    legend.text.align = 1,
    text=element_text(size=12),
    axis.text.x = element_text(angle=0, hjust=0.5, vjust=0.5),
    legend.title.align = 0,
    legend.spacing.y = unit(1, "pt"),
    # plot.title = element_text(hjust=-0.05),
    axis.text.y=element_blank(),
    axis.ticks.y=element_blank(),
    strip.background = element_rect(fill = "white")
  )

# get uncertainty in predictions for panel B

# draw all samples from posteriors
draws_spr <- as_draws_df(m_rt, "^b\\_.*", regex=TRUE)

# function for calculating uncertainty around predictions for a given vector of response times
get_pred_cr_i <- function(rt_i) {
  cat(sprintf("\rCalculating densities %s - %s", min(rt_i), max(rt_i)))
  expand_grid(
    .draw = unique(draws_spr$.draw),
    val_cong = unique(d$cong_dev),
    prop_agree = 1:10/10,
    rt = rt_i
  ) %>%
    left_join(draws_spr, by=".draw") %>%
    mutate(
      val_pred = norm01_manual(prop_agree, min(d$prop_agree), max(d$prop_agree)),
      mu = b_Intercept +
        (val_cong * b_cong_dev) +
        (val_pred * b_pred_norm) +
        (val_cong * val_pred * `b_cong_dev:pred_norm`),
      sigma = b_sigma_Intercept +
        (val_cong * b_sigma_cong_dev) +
        (val_pred * b_sigma_pred_norm) +
        (val_cong * val_pred * `b_sigma_cong_dev:pred_norm`),
      delta = b_ndt_Intercept +
        (val_cong * b_ndt_cong_dev) +
        (val_pred * b_ndt_pred_norm) +
        (val_cong * val_pred * `b_ndt_cong_dev:pred_norm`),
      samp_dens = dshifted_lnorm(
        x = rt,
        meanlog = mu,
        sdlog = exp(sigma),
        shift = exp(delta)
      )
    ) %>%
    group_by(rt, val_cong, val_pred, prop_agree) %>%
    summarise(
      pred_dens = median(samp_dens),
      cr_i_low = hdi(samp_dens, .width=.89)[1],
      cr_i_high = hdi(samp_dens, .width=.89)[2],
      .groups = "drop"
    )
}

# get relative likelihoods (chunked into groups of size 25)
draws_pred_ci <- quantities[quantities>0] %>%
  split(., ceiling(seq_along(.)/25)) %>%
  map_dfr(get_pred_cr_i)

# join panel A and panel B
max_y_uncertainty <- round(max(draws_pred_ci$cr_i_high)+.00005, 5)

pub_panel_A_uncertainty <- pub_panel_A +
  lims(y = c(0, max_y_uncertainty))

pub_panel_B_uncertainty <- draws_pred_ci %>%
  left_join(coding_lookup, by=c("val_cong" = "cong_dev")) %>%
  mutate(
    condition_label = ifelse(condition=="A1", "Congruent", "Incongruent"),
    pred_label = factor(sprintf("%s%%", prop_agree*100), levels = sprintf("%s%%", seq(10, 100, 10)))
  ) %>%
  ggplot(aes(rt, pred_dens, colour = condition_label, fill = condition_label)) +
  geom_ribbon(aes(ymin = cr_i_low, ymax = cr_i_high), alpha=0.4) +
  facet_wrap(vars(pred_label), nrow=2) +
  scale_colour_manual(values = cong_cols) +
  scale_fill_manual(values = cong_cols) +
  guides(fill = guide_legend(override.aes = list(alpha = 0.5))) +
  lims(y = c(0, max_y_uncertainty)) +
  labs(
    x = "Response Time (ms)",
    y = "Predicted Density",
    colour = "Picture-Word Congruency",
    fill = "Picture-Word Congruency",
    tag = "b"
  ) +
  theme(
    legend.position = "bottom",
    legend.key.height = unit(4, "pt"),
    axis.text.y=element_blank(),
    axis.ticks.y=element_blank(),
    axis.text.x = element_text(angle=22.5, hjust=1, vjust=1),
    # legend.key.height = grid::unit(0.1, "lines"),
    # plot.title = element_text(hjust=-0.04),
    strip.background = element_rect(fill = "white"),
    legend.margin = margin()
  )

pub_fig_uncertainty <- plot_grid(pub_panel_A_uncertainty, pub_panel_B_uncertainty, nrow=2, rel_heights=c(2.5, 3.5))

ggsave(file.path("figs", "08_rt_fixed_effects_uncertainty.pdf"), pub_fig_uncertainty, device = "pdf", units = "in", width = 6.5, height=6)
ggsave(file.path("figs", "08_rt_fixed_effects_uncertainty.png"), pub_fig_uncertainty, device = "png", type="cairo", units = "in", width = 6.5, height=6)


# compare priors and posteriors -------------------------------------------

m_rt_prior_samps <- brm(
  formula = f,
  data = d,
  family = shifted_lognormal(),
  prior = priors,
  iter = n_iter,
  warmup = n_warmup,
  chains = n_chains,
  control = list(
    adapt_delta = adapt_delta,
    max_treedepth = max_treedepth
  ),
  inits = replicate(
    n_chains,
    list(b_ndt = as.array(rep(-5, 4))),
    simplify=FALSE
  ),
  sample_prior = "only",
  silent = TRUE,
  cores = n_cores,
  seed = seed,
  thin = 1,
  refresh = 2500
)

rm(draws_spr)
gc()

draws_joined <- bind_rows(
  as_draws_df(m_rt, "^b\\_.*|^sd\\_.*", regex=TRUE) %>%
    select(-.chain, -.iteration, -.draw) %>%
    pivot_longer(cols=everything(), names_to="par", values_to="est") %>%
    mutate(source="posterior"),
  as_draws_df(m_rt_prior_samps, "^b\\_.*|^sd\\_.*", regex=TRUE) %>%
    select(-.chain, -.iteration, -.draw) %>%
    pivot_longer(cols=everything(), names_to="par", values_to="est") %>%
    mutate(source="prior")
) %>%
  mutate(source = factor(source, levels = c("prior", "posterior")))


pl_prior_post_fe_ints <- draws_joined %>%
  filter(grepl("^b\\_", par), grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(recode(
      par,
      b_Intercept = "mu",
      b_sigma_Intercept = "sigma",
      b_ndt_Intercept = "delta"
    ), levels = c("mu", "sigma", "delta"))
  ) %>%
  ggplot(aes(est, "Intercept", colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red")) +
  labs(
    x = NULL,
    y = NULL
  ) +
  theme(legend.position = "none")

pl_prior_post_fe_slopes <- draws_joined %>%
  filter(grepl("^b\\_", par), !grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(case_when(
      grepl("sigma", par, fixed=TRUE) ~ "sigma",
      grepl("ndt", par, fixed=TRUE) ~ "delta",
      TRUE ~ "mu"
    ), levels = c("mu", "sigma", "delta")),
    eff = factor(case_when(
      grepl("cong_dev:pred_norm", par, fixed=TRUE) ~ "Congruency\n* Predictability",
      grepl("cong_dev", par, fixed=TRUE) ~ "Congruency",
      grepl("pred_norm", par, fixed=TRUE) ~ "Predictability"
    ), levels = c("Congruency", "Predictability", "Congruency\n* Predictability"))
  ) %>%
  ggplot(aes(est, reorder(eff, desc(eff)), colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red"), labels = c("Prior", "Posterior")) +
  labs(
    x = "Estimate",
    y = NULL,
    colour = NULL
  ) +
  theme(
    legend.position = "bottom",
    legend.margin = margin(),
    strip.background = element_blank(),
    strip.text.x = element_blank()
  )

pl_prior_post_fe <- plot_grid(pl_prior_post_fe_ints, pl_prior_post_fe_slopes, align="hv", axis="l", ncol=1, rel_heights=c(1.25, 2.85))

ggsave(file.path("figs", "08_rt_prior_post_fixed_effects.pdf"), pl_prior_post_fe, width=6.5, height=3.5)
ggsave(file.path("figs", "08_rt_prior_post_fixed_effects.png"), pl_prior_post_fe, width=6.5, height=3.5, device="png", type="cairo")


# random effects plot

# subject random effects SDs
pl_prior_post_re_subj_ints <- draws_joined %>%
  filter(grepl("^sd\\_subj\\_id", par), grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(case_when(
      grepl("sigma", par, fixed=TRUE) ~ "sigma",
      grepl("ndt", par, fixed=TRUE) ~ "delta",
      TRUE ~ "mu"
    ), levels = c("mu", "sigma", "delta"))
  ) %>%
  ggplot(aes(est, "Intercept", colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red")) +
  labs(
    x = NULL,
    y = NULL,
    title = "Participant Random Effects SDs",
    tag = "a"
  ) +
  theme(legend.position = "none")

pl_prior_post_re_subj_slopes <- draws_joined %>%
  filter(grepl("^sd\\_subj\\_id", par), !grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(case_when(
      grepl("sigma", par, fixed=TRUE) ~ "sigma",
      grepl("ndt", par, fixed=TRUE) ~ "delta",
      TRUE ~ "mu"
    ), levels = c("mu", "sigma", "delta")),
    eff = factor(case_when(
      grepl("cong_dev:pred_norm", par, fixed=TRUE) ~ "Congruency\n* Predictability",
      grepl("cong_dev", par, fixed=TRUE) ~ "Congruency",
      grepl("pred_norm", par, fixed=TRUE) ~ "Predictability"
    ), levels = c("Congruency", "Predictability", "Congruency\n* Predictability"))
  ) %>%
  ggplot(aes(est, reorder(eff, desc(eff)), colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red"), labels = c("Prior", "Posterior")) +
  labs(
    x = NULL,
    y = NULL,
    colour = NULL
  ) +
  theme(
    legend.position = "none",
    strip.background = element_blank(),
    strip.text.x = element_blank()
  )

# image random effects SDs
pl_prior_post_re_image_ints <- draws_joined %>%
  filter(grepl("^sd\\_image", par), grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(case_when(
      grepl("sigma", par, fixed=TRUE) ~ "sigma",
      grepl("ndt", par, fixed=TRUE) ~ "delta",
      TRUE ~ "mu"
    ), levels = c("mu", "sigma", "delta"))
  ) %>%
  ggplot(aes(est, "Intercept", colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red")) +
  labs(
    x = NULL,
    y = NULL,
    title = "Image Random Effects SDs",
    tag = "b"
  ) +
  theme(legend.position = "none")

pl_prior_post_re_image_slopes <- draws_joined %>%
  filter(grepl("^sd\\_image", par), !grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(case_when(
      grepl("sigma", par, fixed=TRUE) ~ "sigma",
      grepl("ndt", par, fixed=TRUE) ~ "delta",
      TRUE ~ "mu"
    ), levels = c("mu", "sigma", "delta"))
  ) %>%
  ggplot(aes(est, "Congruency", colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red"), labels = c("Prior", "Posterior")) +
  labs(
    x = NULL,
    y = NULL,
    colour = NULL
  ) +
  theme(
    legend.position = "none",
    strip.background = element_blank(),
    strip.text.x = element_blank()
  )
  
# word random effects SDs
pl_prior_post_re_string_ints <- draws_joined %>%
  filter(grepl("^sd\\_string", par), grepl("Intercept", par, fixed=TRUE)) %>%
  mutate(
    par_lab = factor(case_when(
      grepl("sigma", par, fixed=TRUE) ~ "sigma",
      grepl("ndt", par, fixed=TRUE) ~ "delta",
      TRUE ~ "mu"
    ), levels = c("mu", "sigma", "delta"))
  ) %>%
  ggplot(aes(est, "Intercept", colour=source)) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89, position=position_dodge(width=-0.4)) +
  facet_wrap(vars(par_lab), scales = "free_x", labeller = label_parsed) +
  scale_y_discrete(expand = expansion(0.1, 0)) +
  scale_colour_manual(values = c("black", "red"), labels=c("Prior", "Posterior")) +
  labs(
    x = "Estimate",
    y = NULL,
    title = "Word Random Effects SDs",
    tag = "c",
    colour = NULL
  ) +
  theme(legend.position = "bottom", legend.margin = margin())

# join random effects SDs plots
pl_prior_post_re <- plot_grid(
  pl_prior_post_re_subj_ints, pl_prior_post_re_subj_slopes,
  pl_prior_post_re_image_ints, pl_prior_post_re_image_slopes,
  pl_prior_post_re_string_ints,
  align="hv", axis="l", ncol=1, rel_heights=c(0.9, 1.2, 0.9, 0.5, 1.255)
)

ggsave(file.path("figs", "08_rt_prior_post_random_effects.pdf"), pl_prior_post_re, width=6.5, height=7.5)
ggsave(file.path("figs", "08_rt_prior_post_random_effects.png"), pl_prior_post_re, width=6.5, height=7.5, device="png", type="cairo")