# this script is the same as analyse_03 but uses the data from the word vs noise maximal electrode instead of word vs consonants

library(lme4)
library(brms)
library(ggdist)
library(dplyr)
library(purrr)
library(readr)
library(tidyr)
library(ggplot2)
library(scales)
library(patchwork)
library(ggnewscale)

ggplot2::theme_set(ggplot2::theme_classic() + theme(strip.background = element_rect(fill = "white")))

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

cong_cols_light <- sapply(cong_cols, function(x) {
  x_rgb <- as.numeric(col2rgb(x))
  x_li_rgb <- round(rowMeans(cbind(x_rgb*1.25, c(255, 255, 255)*0.75)))
  x_li <- rgb(x_li_rgb[1], x_li_rgb[2], x_li_rgb[3], maxColorValue=255)
  x_li
}, USE.NAMES=FALSE)

# function to normalise between 0 and 1
norm01 <- function(x, ...) (x-min(x, ...))/(max(x, ...)-min(x, ...))

# 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)

# import the max electrode data from the preprocessing, and set up the variables for the model
d <- list.files("max_elec_data", pattern=".+\\.csv$", full.names=TRUE) %>%
  map_dfr(function(f) {
    read_csv(
      f,
      col_types=cols(
        date = col_date(format = "%d/%m/%Y"),
        trial_save_time = col_time(format = "%H:%M:%S"),
        subj_id = col_character(),
        stim_grp = col_integer(),
        resp_grp = col_integer(),
        sex = col_character(),
        trl_nr = col_integer(),
        item_nr = col_integer(),
        condition = col_character(),
        eeg_trigg = col_integer(),
        image = col_character(),
        string = col_character(),
        corr_ans = col_character(),
        resp = col_character(),
        acc = col_integer(),
        fix1_jitt_flip = col_integer(),
        fix2_jitt_flip = col_integer(),
        event = col_integer(),
        is_block_start = col_logical(),
        is_block_end = col_logical(),
        .default = col_double()
      )
    ) %>%
      select(
        subj_id, stim_grp, resp_grp, item_nr, condition,
        image, string, acc, rt, uV, uV_noise_elec
      )
  }) %>%
  left_join(stim, by=c("image" = "filename")) %>%
  mutate(
    cong_dev = as.numeric(scale(ifelse(condition=="A2", 0, 1), center=TRUE, scale=FALSE)),
    cong_dum_incong = as.numeric(condition=="A1"),  # 0 is at incongruent
    cong_dum_cong = as.numeric(condition=="A2"),  # 0 is at congruent
    prop_agree = perc_name_agree/100,
    pred_norm = norm01(prop_agree, na.rm=TRUE),
    pred_norm_max = pred_norm - 1
  )

# fit the model ----
m <- lmer(
  uV_noise_elec ~ cong_dev * pred_norm +
    (cong_dev * pred_norm | subj_id) +
    (cong_dev | image) +
    (1 | string),
  REML=FALSE,
  control = lmerControl(optimizer="bobyqa"),
  data=d
)

# the effect of the interaction
m_no_interact <- update(m, ~. -cong_dev:pred_norm)
eff_interact <- anova(m, m_no_interact)

# decompose interaction by congruency ----

# effect of predictability in incongruent trials
m_incong <- lmer(
  uV_noise_elec ~ cong_dum_incong * pred_norm +
    (cong_dum_incong * pred_norm | subj_id) +
    (cong_dum_incong | image) +
    (1 | string),
  REML=FALSE,
  control = lmerControl(optimizer="bobyqa"),
  data=d
)

m_incong_no_pred <- update(m_incong, ~. -pred_norm)
eff_pred_incong <- anova(m_incong, m_incong_no_pred)

# effect of predictability in congruent trials
m_cong <- lmer(
  uV_noise_elec ~ cong_dum_cong * pred_norm +
    (cong_dum_cong * pred_norm | subj_id) +
    (cong_dum_cong | image) +
    (1 | string),
  REML=FALSE,
  control = lmerControl(optimizer="bobyqa"),
  data=d
)

m_cong_no_pred <- update(m_cong, ~. -pred_norm)
eff_pred_cong <- anova(m_cong, m_cong_no_pred)

# decompose interaction by predictability ----

# effect of congruency at minimum level of predictability is in main model

# effect of congruency at maximum predictability
m_maxpred <- lmer(
  uV_noise_elec ~ cong_dev * pred_norm_max +
    (cong_dev * pred_norm_max | subj_id) +
    (cong_dev | image) +
    (1 | string),
  REML=FALSE,
  control = lmerControl(optimizer="bobyqa"),
  data=d
)

m_maxpred_no_cong <- update(m_maxpred, ~. -cong_dev)
eff_cong_maxpred <- anova(m_maxpred, m_maxpred_no_cong)

# plot relationship ----

d_cells_pred <- tibble(
  pred_norm = rep(range(d$pred_norm), 2),
  perc_name_agree = rep(range(d$perc_name_agree), 2),
  cong_dev = rep(unique(d$cong_dev), each=2),
  condition = ifelse(cong_dev==min(cong_dev), "A2", "A1")
) %>%
  mutate(pred_uV = predict(m, newdata=., re.form=~0))

# get prediction intervals (no random slopes for feasibility)
m_no_rand_s <- lmer(
  uV_noise_elec ~ cong_dev * pred_norm +
    (1 | subj_id) +
    (1 | image) +
    (1 | string),
  REML=FALSE,
  control = lmerControl(optimizer="bobyqa"),
  data=d
)

d_preds_boot <- expand_grid(
  cong_dev = unique(d$cong_dev),
  perc_name_agree = seq(min(d$perc_name_agree), max(d$perc_name_agree), 0.01)
) %>%
  mutate(
    prop_agree = perc_name_agree/100,
    pred_norm = norm01(prop_agree, na.rm=TRUE),
    condition = ifelse(cong_dev==min(cong_dev), "A2", "A1")
  )

boots <- bootMer(m_no_rand_s, nsim=5000, FUN=function(m_i) {
  predict(m_i, newdata=d_preds_boot, re.form=~0)
}, seed=3101, .progress="txt")

ci_norm <- function(x, level=c(0.025, 0.975)) {
  qnorm(p=level, mean=mean(x), sd=sd(x))
}

d_preds_boot <- d_preds_boot %>%
  mutate(
    pred_int_lwr = apply(boots$t, 2, ci_norm, level=.025),
    pred_int_upr = apply(boots$t, 2, ci_norm, level=.975)
  )

lmm_plot_scatter <- d %>%
  ggplot(aes(perc_name_agree, uV_noise_elec, colour = condition)) +
  geom_point(shape=1, show.legend=FALSE, alpha=0.25) +
  scale_colour_manual(name = NA, labels = NULL, values = cong_cols_light) +
  new_scale_colour() +
  geom_line(aes(y=pred_uV, colour = condition), data=d_cells_pred, linewidth=1.75) +
  scale_colour_manual(name = "Picture-Word Congruency", labels = c("Congruent", "Incongruent"), values = cong_cols) +
  scale_y_continuous(breaks=scales::extended_breaks(n=6)) +
  labs(x = "Predictability (%)", y = "N1 Amplitude (µV)", tag="a")

lmm_plot_lines <- ggplot() +
  geom_ribbon(aes(x=perc_name_agree, ymin=pred_int_lwr, ymax=pred_int_upr, colour=condition), data=d_preds_boot, fill=NA, linetype="dashed", show.legend = FALSE) +
  geom_line(aes(x=perc_name_agree, y=pred_uV, colour=condition), data=d_cells_pred, linewidth=1.75) +
  scale_colour_manual(name = "Picture-Word Congruency", labels = c("Congruent", "Incongruent"), values = cong_cols) +
  scale_y_continuous(breaks=scales::extended_breaks(n=6)) +
  labs(x = "Predictability (%)", y = "N1 Amplitude (µV)", tag="b")

lmm_plot <- (lmm_plot_scatter | lmm_plot_lines) +
  plot_layout(guides = "collect") &
  theme(legend.position = "bottom", plot.background = element_blank()) &
  scale_x_continuous(expand=c(0, 0))

ggsave(file.path("figs", "04_lmm_summary_plot_noise_elec.png"), lmm_plot, device="png", type="cairo", width=6.5, height=3.75, dpi=600)

ggsave(file.path("figs", "04_lmm_summary_plot_noise_elec.pdf"), lmm_plot, device="pdf", width=6.5, height=3.75)

# evidences ratio for effect -----------------------------------------------

p <- c(
  set_prior("normal(-5, 10)", class="Intercept"),
  set_prior("normal(0, 5)", class="b", coef="cong_dev"),
  set_prior("normal(0, 5)", class="b", coef="pred_norm"),
  set_prior("normal(0, 5)", class="b", coef="cong_dev:pred_norm"),
  set_prior("lkj_corr_cholesky(1)", class="L")
)

m_b <- brm(
  uV_noise_elec ~ cong_dev * pred_norm +
    (cong_dev * pred_norm | subj_id) +
    (cong_dev | image) +
    (1 | string),
  data = d,
  prior = p,
  chains = 5,
  cores = 5,
  iter = 5000,
  control = list(
    adapt_delta = 0.8,
    max_treedepth = 10
  ),
  refresh = 25,
  seed = 420
)

m_b_h <- hypothesis(m_b, "cong_dev:pred_norm>0")

m_b_h_hdi <- as_draws_df(m_b, variable="b_cong_dev:pred_norm", regex=FALSE) %>%
  rename(int_est = `b_cong_dev:pred_norm`) %>%
  median_hdi(int_est, .width=0.89)

m_b_pl <- as_draws_df(m_b, variable="b_cong_dev:pred_norm", regex=FALSE) %>%
  rename(int_est = `b_cong_dev:pred_norm`) %>%
  ggplot(aes(int_est, fill=after_stat(x>0))) +
  stat_slab(slab_colour="black", slab_size=0.4) +
  stat_pointinterval(
    .width=0.89, point_interval="median_hdi",
    point_size=1, size=1, show.legend = FALSE
  ) +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_text(
    aes(x = x, y = y),
    data = tibble(x=-3.9, y=1),
    label = deparse(bquote(BF["01"]==.(round(1/m_b_h$hypothesis$Evid.Ratio, 2)))),
    parse = TRUE,
    hjust = 0,
    vjust = 1,
    size = 4.5
  ) +
  scale_y_continuous(expand = expansion(0, 0.04), limits = c(0, NA)) +
  coord_cartesian() +
  labs(
    x = "Congruency-Predictability Interaction (µV)",
    y = "Posterior Density"
  ) +
  theme(legend.position = "none") +
  scale_fill_manual(
    values = c("grey90", "#8B0000"),
    labels = c(bquote(H["0"]), bquote(H["1"]))
  )

ggsave(file.path("figs", "04_post.png"), width=3.5, height=1.75, device="png", type="cairo")
ggsave(file.path("figs", "04_post.pdf"), width=3.5, height=1.75, device="pdf")


lmm_plot_post <- wrap_plots(
  lmm_plot,
  (
    m_b_pl +
      labs(
        fill = NULL,
        tag = "c",
        x = "Congruency-Predictability\nInteraction"
      ) +
      theme(
        plot.background = element_blank(),
        legend.position = "bottom"
      )
  )
) +
  plot_layout(widths = c(2, 1))

ggsave(file.path("figs", "04_lmm_and_post.pdf"), lmm_plot_post, width=6.5, height=3, device="pdf")