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

library(brms)
library(ggdist)

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

library(RColorBrewer)
eff_cols <- c("#000000", brewer.pal(3, "Set1")[1:2], brewer.pal(5, "Set1")[5])

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)

# 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(rt <= 1500) %>%
  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(
  set_prior("normal(4, 1)", class="b", coef="Intercept")
)

n_cores <- 7

seed <- 3101
n_iter <- 10000
n_warmup <- 5000
adapt_delta <- 0.9
max_treedepth <- 10
n_chains <- 5
refresh <- 100

f <- brmsformula(
  acc ~ 0 + Intercept + cong_dev * pred_norm +
    (cong_dev * pred_norm | subj_id) +
    (cong_dev | image) +
    (1 | string)
)

m_acc <- brm(
  formula = f,
  data = d,
  family = bernoulli("logit"),
  prior = priors,
  iter = n_iter,
  warmup = n_warmup,
  chains = n_chains,
  control = list(
    adapt_delta = adapt_delta,
    max_treedepth = max_treedepth
  ),
  sample_prior = "no",
  silent = TRUE,
  cores = n_cores,
  seed = seed,
  thin = 1,
  file = file.path("mods", "m_acc.rds"),
  refresh = refresh
)

m_draws <- as_draws_df(m_acc, variable="^b\\_", regex=TRUE) %>%
  select(-.chain, -.iteration, -.draw)

m_draws_long <- m_draws %>%
  pivot_longer(cols=everything(), names_to="par", values_to="est") %>%
  mutate(par_lab = factor(recode(
    par,
    b_Intercept = "Intercept",
    b_cong_dev = "Congruency",
    b_pred_norm = "Predictability",
    `b_cong_dev:pred_norm` = "Congruency\n* Predictability"
  ), levels = c("Intercept", "Congruency", "Predictability", "Congruency\n* Predictability")))

pl_intercept <- m_draws_long %>%
  filter(par == "b_Intercept") %>%
  ggplot(aes(est, "Intercept")) +
  stat_pointinterval(point_interval = "median_hdi", .width=.89) +
  labs(x = NULL, y = NULL, tag = "a")

pl_slopes <- m_draws_long %>%
  filter(par != "b_Intercept") %>%
  mutate(par_lab = fct_rev(par_lab)) %>%
  ggplot(aes(est, par_lab)) +
  geom_vline(xintercept=0, linetype="dashed") +
  stat_pointinterval(point_interval = "median_hdi", .width=.89) +
  labs(x = "Estimate (Logits)", y = NULL) +
  theme(legend.position = "none")

pl_coefs <- (pl_intercept / pl_slopes) +
  plot_layout(heights = c(1, 5))

m_preds <- seq(0.1, 1, 0.001) %>%
  split(., ceiling(seq_along(.)/100)) %>%
  map_dfr(function(p) {
    m_draws %>%
      expand_grid(
        prop_agree = p,
        condition = c("Congruent", "Incongruent")
      ) %>%
      mutate(
        pred_norm = norm01_manual(prop_agree, min_x=min(d$prop_agree), max_x=max(d$prop_agree)),
        cong_dev = as.numeric(scale(if_else(condition == "Congruent", 1, 0), center = TRUE, scale = FALSE)),
        pred_logit = b_Intercept +
          cong_dev * b_cong_dev +
          pred_norm * b_pred_norm +
          cong_dev * pred_norm * `b_cong_dev:pred_norm`,
        pred_odds = exp(pred_logit),
        prob = pred_odds / (1 + pred_odds)
      ) %>%
      group_by(condition, prop_agree) %>%
      median_hdi(prob, .width=0.89)
  }) %>%
  mutate(
    perc_name_agree = prop_agree * 100,
    condition = factor(condition, levels = c("Congruent", "Incongruent"))
  )

pl_preds <- m_preds %>%
  ggplot(aes(perc_name_agree, prob, colour=condition, fill=condition)) +
  geom_line(size=0.5) +
  geom_ribbon(aes(ymin=.lower, ymax=.upper), size=0.001, alpha=0.25) +
  scale_colour_manual(values = cong_cols, guide=guide_legend(nrow=1)) +
  scale_fill_manual(values = cong_cols, guide=guide_legend(nrow=1)) +
  labs(
    x = "Predictability (%)",
    y = "Probability of Correct Response",
    colour = "Picture-Word Congruency",
    fill = "Picture-Word Congruency",
    tag = "b"
  ) +
  scale_x_continuous(expand = expansion()) +
  scale_y_continuous(expand = expansion(), limits=c(NA, 1)) +
  theme(
    legend.position = c(0.95, 0.05),
    legend.justification = c(1, 0),
    legend.key.height = unit(4, "pt")
  )

res_pl <- (pl_coefs | pl_preds) +
  plot_layout(widths = c(0.55, 1))

ggsave(file.path("figs", "09_pictureword_acc_res.pdf"), res_pl, width=6.5, height=3)
ggsave(file.path("figs", "09_pictureword_acc_res.png"), device="png", type="cairo", res_pl, width=6.5, height=3)