library(lme4)
library(dplyr)
library(tidyr)
library(purrr)
library(readr)
library(ggplot2)
library(ggtext)
library(cowplot)
library(patchwork)
library(parallel)
ggplot2::theme_set(ggplot2::theme_classic())
cong_cols <- c("#E69F00", "#009E73")
n_cores <- 14

library(RColorBrewer)
eff_cols <- c("#000000", brewer.pal(3, "Set1")[1:2], brewer.pal(5, "Set1")[5])
cong_cols <- c("#E69F00", "#009E73")
cong_cols_alpha_0.4_white <- c("#F5D899", "#99D8C7")

roi <- c("TP7", "CP5", "P5", "P7", "P9", "PO3", "PO7", "O1")
time_lims <- c(-250, 650)

# 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
get_sample_data <- function(ch_i = NA) {
  list.files("sample_data", pattern=".+\\.csv$", full.names=TRUE) %>%
    map_dfr(function(f) {
      x <- read_csv(
        f,
        col_types=cols(
          subj_id = col_character(),
          stim_grp = col_integer(),
          resp_grp = col_integer(),
          item_nr = col_integer(),
          ch_name = col_character(),
          time = col_double(),
          uV = col_double(),
          condition = col_character(),
          image = col_character(),
          string = col_character(),
          .default = col_double()
        ),
        progress = FALSE
      ) |>
        select(-stim_grp, -resp_grp, -item_nr) |>
        filter(time > time_lims[[1]] & time < time_lims[[2]])
      if (all(!is.na(ch_i))) x <- filter(x, ch_name %in% ch_i)
      x
    }) %>%
    left_join(stim, by=c("image" = "filename")) %>%
    mutate(
      prop_agree = perc_name_agree/100,
      pred_norm = norm01(prop_agree, na.rm=TRUE),
      pred_norm_max = pred_norm - 1
    ) %>%
    select(-perc_name_agree)
}

unique_ids <- list.files("sample_data", pattern=".+\\.csv$", full.names=TRUE) %>%
  first() %>%
  read_csv() |>
  filter(time > time_lims[[1]] & time < time_lims[[2]]) |>
  select(ch_name, time) %>%
  as.list() %>%
  lapply(unique)

target_time_ids <- unique_ids$time

# get models for each timepoint, for each electrode, and extract fixed effects
d_i_list <- get_sample_data(roi) |>
  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
  ) |>
  group_split(time)

gc()

# # sanity check
# sc_pl <- d_i_list |>
#   reduce(bind_rows) |>
#   mutate(pred_bin = factor(case_when(
#     pred_norm <.33 ~ "low",
#     pred_norm <.66 ~ "medium",
#     pred_norm < Inf ~ "high"
#   ), levels = c("low", "medium", "high"))) |>
#   group_by(time, condition, pred_bin, ch_name) |>
#   summarise(mean_uV = mean(uV)) |>
#   ungroup() |>
#   pivot_wider(names_from=condition, values_from=mean_uV) |>
#   mutate(cong_minus_incong = A1-A2) |>
#   ggplot(aes(time, cong_minus_incong, colour=ch_name)) +
#   geom_line() +
#   facet_wrap(vars(pred_bin))

message("Fitting models in parallel to OT channels in picture-word experiment")

# d_i_list <- get_sample_data(roi) %>%
#   mutate(cong_dev = as.numeric(scale(ifelse(condition=="A2", 0, 1), center=TRUE, scale=FALSE))) %>%
#   filter(time <= 655) %>%
#   arrange(time) %>%
#   split(time)
# 
# gc()

cl <- makeCluster(n_cores)
cl_packages <- clusterEvalQ(cl, {
  library(dplyr)
  library(lme4)
})

fe_res <- parLapply(cl, d_i_list, function(d_i) {
  # m <- lme4::lmer(
  #   uV ~ cong_dev * pred_norm +
  #     (cong_dev * pred_norm | subj_id) +
  #     (cong_dev * pred_norm | ch_name:subj_id) +
  #     (cong_dev | image) +
  #     (1 | string),
  #   REML=FALSE,
  #   control = lmerControl(optimizer="bobyqa"),
  #   data=d_i
  # )
  m <- lme4::lmer(
    uV ~ cong_dev * pred_norm +
      (1 | subj_id) +
      (1 | ch_name:subj_id) +
      (1 | image) +
      (1 | string),
    REML=FALSE,
    control = lmerControl(optimizer="bobyqa"),
    data=d_i
  )
  
  m %>%
    summary() %>%
    with(coefficients) %>%
    as_tibble(rownames = "fe")
}) %>%
  reduce(bind_rows) %>%
  mutate(time = rep(target_time_ids, each=4))

stopCluster(cl)

fe_res_tidy <- fe_res %>%
  mutate(
    fe_lab = factor(recode(
      fe,
      `(Intercept)` = "Intercept",
      cong_dev = "Congruency",
      pred_norm = "Predictability",
      `cong_dev:pred_norm` = "Congruency * Predictability",
    ), levels = c("Intercept", "Congruency", "Predictability", "Congruency * Predictability")),
    fe_lab_newline = fe_lab,
    bound_lower = Estimate - 1.96 * `Std. Error`,
    bound_upper = Estimate + 1.96 * `Std. Error`
  )

levels(fe_res_tidy$fe_lab) <- c("Intercept"="Intercept", "Congruency"="Congruency", "Predictability"="Predictability", "Congruency * Predictability"="Congruency × Predictability")

levels(fe_res_tidy$fe_lab_newline) <- c("Intercept"="Intercept", "Congruency"="Congruency", "Predictability"="Predictability", "Congruency * Predictability"="Congruency<br>× Predictability")

overlay_dat <- fe_res_tidy %>%
  filter(fe_lab=="Intercept") %>%
  select(-fe_lab, -fe_lab_newline) %>%
  expand_grid(fe_lab = unique(fe_res_tidy$fe_lab)) %>%
  mutate(
    Estimate = ifelse(fe_lab=="Intercept", NA, Estimate),
    fe_lab_newline = fe_lab
  )

levels(overlay_dat$fe_lab) <- c("Intercept"="Intercept", "Congruency"="Congruency", "Predictability"="Predictability", "Congruency * Predictability"="Congruency × Predictability")

levels(overlay_dat$fe_lab_newline) <- c("Intercept"="Intercept", "Congruency"="Congruency", "Predictability"="Predictability", "Congruency * Predictability"="Congruency<br>× Predictability")

ylims <- round(c(min(fe_res_tidy$bound_lower), max(fe_res_tidy$bound_upper)))

pl_roi <- fe_res_tidy %>%
  ggplot(aes(time, Estimate, ymin=bound_lower, ymax=bound_upper, group=fe_lab)) +
  geom_line(colour="black", data=overlay_dat, alpha=0.6) +
  geom_ribbon(alpha=0.5, linewidth=0.1, fill="dodgerblue") +
  geom_line() +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  facet_wrap(vars(fe_lab), nrow=2) +
  labs(x = "Time (ms)", y = "Fixed Effect Estimate (µV)", tag="a") +
  # scale_colour_manual(values = eff_cols) +
  # scale_fill_manual(values = eff_cols) +
  scale_x_continuous(expand = c(0, 0)) +
  scale_y_continuous(limits=ylims) +
  theme(
    legend.position = "none",
    strip.background = element_blank(),
    legend.background = element_blank(),
    axis.line.x = element_blank(),
    panel.spacing.x = unit(12, "pt"),
    strip.text.x = ggtext::element_markdown()
  ) +
  coord_cartesian(xlim=c(-250, time_lims[2]))

pl_roi_1row <- fe_res_tidy %>%
  ggplot(aes(time, Estimate, ymin=bound_lower, ymax=bound_upper, group=fe_lab)) +
  geom_line(colour="darkgrey", data=overlay_dat) +
  geom_ribbon(alpha=0.5, linewidth=0.1, fill="dodgerblue") +
  geom_line() +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  facet_wrap(vars(fe_lab_newline), nrow=1) +
  labs(x = "Time (ms)", y = "Effect Estimate (µV)", tag="a") +
  # scale_colour_manual(values = eff_cols) +
  # scale_fill_manual(values = eff_cols) +
  scale_x_continuous(expand = c(0, 0)) +
  scale_y_continuous(limits=ylims) +
  theme(
    legend.position = "none",
    strip.background = element_blank(),
    legend.background = element_blank(),
    axis.line.x = element_blank(),
    panel.spacing.x = unit(12, "pt"),
    strip.text.x = ggtext::element_markdown(),
    strip.clip = "off"
  ) +
  coord_cartesian(xlim=c(-250, time_lims[2]))

preds <- expand_grid(
  prop_agree = seq(0.1, 1, 0.1),
  condition = factor(c("Picture-Congruent", "Picture-Incongruent"), levels=c("Picture-Congruent", "Picture-Incongruent"))
) %>%
  rowwise() %>%
  mutate(pred_norm = ifelse(prop_agree==1, 1, norm01(c(0.07, prop_agree, 1))[2])) %>%
  expand_grid(time = target_time_ids) %>%
  ungroup() %>%
  mutate(
    perc_name_agree = prop_agree*100,
    cong_dev = as.numeric(scale(ifelse(condition=="Picture-Incongruent", 0, 1), center=TRUE, scale=FALSE))
  ) %>%
  left_join(
    fe_res %>%
      select(time, fe, Estimate) %>%
      pivot_wider(names_from=fe, values_from=Estimate, names_prefix="fe_"),
    by = "time"
  ) %>%
  mutate(
    pred_uV = `fe_(Intercept)` +
      cong_dev * fe_cong_dev +
      pred_norm * fe_pred_norm +
      cong_dev * pred_norm * `fe_cong_dev:pred_norm`
  )

pl_preds <- preds %>%
  ggplot(aes(time, pred_uV, colour=prop_agree)) +
  geom_line(aes(group = as.factor(prop_agree))) +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  scale_colour_continuous(
    type="viridis", breaks=sort(unique(preds$prop_agree)),
    labels=sprintf("%s%%", sort(unique(preds$prop_agree))*100),
    # guide=guide_colourbar(barheight = 7.6, barwidth = 1)
    guide = guide_legend(override.aes = list(linewidth=1), reverse = TRUE)
  ) +
  scale_x_continuous(expand=c(0, 0)) +
  scale_y_continuous(limits=ylims) +
  facet_wrap(vars(condition), nrow=1) +
  labs(
    x = "Time (ms)",
    y = "Predicted Amplitude (µV)",
    colour = "Predictability",
    tag = "b"
  ) +
  theme(
    legend.position = "right",
    legend.key.height = unit(11, "pt"),
    legend.text.align = 1,
    strip.background = element_blank(),
    plot.background = element_blank(),
    axis.line.x = element_blank(),
    panel.spacing.x = unit(12, "pt")
  ) +
  coord_cartesian(xlim=c(-250, time_lims[2]))

pl <- plot_grid(pl_roi, pl_preds, ncol=1, rel_heights = c(1, 0.8), align="v", axis="l")

ggsave(file.path("figs", "sample_level", "picture_word", "roi.pdf"), pl, width=6.5, height=6, device="pdf")
ggsave(file.path("figs", "sample_level", "picture_word", "roi.png"), pl, width=6.5, height=6, device="png", type="cairo")

pl_preds_bypred <- preds %>%
  mutate(
    perc_lab = factor(
      sprintf("%s%%", perc_name_agree),
      levels = sprintf("%s%%", sort(unique(perc_name_agree)))
    ),
    cong_lab = ifelse(condition=="Picture-Congruent", "Congruent", "Incongruent")
  ) %>%
  ggplot(aes(time, pred_uV, colour=cong_lab)) +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  geom_line() +
  scale_colour_manual(
    values = cong_cols,
    guide=guide_legend(nrow=1, override.aes = list(linewidth=1))
  ) +
  scale_x_continuous(expand=c(0, 0)) +
  scale_y_continuous(limits=ylims) +
  facet_wrap(vars(perc_lab), nrow=3) +
  labs(
    x = "Time (ms)",
    y = "Predicted Amplitude (µV)",
    colour = "Picture-Word Congruency"
  ) +
  theme(
    legend.position = c(0.925, -0.025),
    legend.justification = c(1, 0),
    legend.margin = margin(),
    legend.key.height = unit(12, "pt"),
    strip.background = element_blank(),
    plot.background = element_blank(),
    axis.line.x = element_blank(),
    panel.spacing.x = unit(12, "pt")
  ) +
  coord_cartesian(xlim=c(-250, time_lims[2]))

pl_plus_C <- plot_grid(
  pl_roi_1row + theme(strip.text = element_text(margin=margin(t=2.5, b=2.5, unit="pt"))),
  pl_preds +
    theme(
      strip.text = element_text(margin=margin(t=2.5, b=2.5, unit="pt")),
      legend.key.height = unit(9.5, "pt")
    ),
  pl_preds_bypred +
    labs(tag="c") +
    theme(strip.text = element_text(margin=margin(t=2.5, b=2.5, unit="pt"))),
  ncol=1, rel_heights = c(0.85, 1.2, 1.5), align="v", axis="l"
)


ggsave(file.path("figs", "sample_level", "picture_word", "roi_3panels.pdf"), pl_plus_C, width=6.5, height=7, device="pdf")
ggsave(file.path("figs", "sample_level", "picture_word", "roi_3panels.png"), pl_plus_C, width=6.5, height=7, device="png", type="cairo")


# simple effects models ---------------------------------------------------

cl <- makeCluster(n_cores)
cl_packages <- clusterEvalQ(cl, {
  library(dplyr)
  library(lme4)
})

fe_res_cong <- parLapply(cl, d_i_list, function(d_i) {
  m <- lme4::lmer(
    uV ~ cong_dum_cong * pred_norm +
      (1 | subj_id) +
      (1 | ch_name:subj_id) +
      (1 | image) +
      (1 | string),
    REML=FALSE,
    control = lmerControl(optimizer="bobyqa"),
    data=d_i
  )
  
  m %>%
    summary() %>%
    with(coefficients) %>%
    as_tibble(rownames = "fe")
}) %>%
  reduce(bind_rows) %>%
  mutate(time = rep(target_time_ids, each=4))

fe_res_incong <- parLapply(cl, d_i_list, function(d_i) {
  m <- lme4::lmer(
    uV ~ cong_dum_incong * pred_norm +
      (1 | subj_id) +
      (1 | ch_name:subj_id) +
      (1 | image) +
      (1 | string),
    REML=FALSE,
    control = lmerControl(optimizer="bobyqa"),
    data=d_i
  )
  
  m %>%
    summary() %>%
    with(coefficients) %>%
    as_tibble(rownames = "fe")
}) %>%
  reduce(bind_rows) %>%
  mutate(time = rep(target_time_ids, each=4))

stopCluster(cl)


fe_res_tidy_se <- bind_rows(
  fe_res_cong %>%
    mutate(
      cong_level = "Congruent",
      fe_lab = factor(recode(
        fe,
        `(Intercept)` = "Intercept",
        cong_dum_cong = "Congruency",
        pred_norm = "Predictability",
        `cong_dum_cong:pred_norm` = "Congruency * Predictability"
      ), levels = c("Intercept", "Congruency", "Predictability", "Congruency * Predictability")),
      fe_lab_newline = factor(fe_lab, labels = c("Intercept", "Congruency", "Predictability", "Congruency\n* Predictability")),
      bound_lower = Estimate - 1.96 * `Std. Error`,
      bound_upper = Estimate + 1.96 * `Std. Error`
    ),
  fe_res_incong %>%
    mutate(
      cong_level = "Incongruent",
      fe_lab = factor(recode(
        fe,
        `(Intercept)` = "Intercept",
        cong_dum_incong = "Congruency",
        pred_norm = "Predictability",
        `cong_dum_incong:pred_norm` = "Congruency * Predictability"
      ), levels = c("Intercept", "Congruency", "Predictability", "Congruency * Predictability")),
      fe_lab_newline = factor(fe_lab, labels = c("Intercept", "Congruency", "Predictability", "Congruency\n* Predictability")),
      bound_lower = Estimate - 1.96 * `Std. Error`,
      bound_upper = Estimate + 1.96 * `Std. Error`
    )
)

fe_res_tidy_se_pred <- filter(fe_res_tidy_se, fe_lab=="Predictability")

fe_se_pl_alpha <- 0.4

fe_se_pl <- fe_res_tidy_se_pred %>%
  ggplot(aes(time, Estimate, ymin=bound_lower, ymax=bound_upper, colour=cong_level, fill=cong_level)) +
  geom_ribbon(alpha=fe_se_pl_alpha, linewidth=0.25, data=filter(fe_res_tidy_se_pred, cong_level=="Congruent")) +
  geom_line(data=filter(fe_res_tidy_se_pred, cong_level=="Congruent")) +
  geom_ribbon(alpha=fe_se_pl_alpha, linewidth=0.25, data=filter(fe_res_tidy_se_pred, cong_level=="Incongruent")) +
  geom_line(data=filter(fe_res_tidy_se_pred, cong_level=="Incongruent")) +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  scale_colour_manual(values = cong_cols) +
  scale_fill_manual(values = cong_cols) +
  scale_x_continuous(expand=c(0, 0), breaks = seq(-200, time_lims[2], 100), limits=c(-250, time_lims[2])) +
  scale_y_continuous(breaks=scales::breaks_pretty(n=6)) +
  guides(
    colour = guide_legend(nrow=1),
    fill = guide_legend(nrow=1, override.aes = list(alpha=1, fill=cong_cols_alpha_0.4_white))
  ) +
  labs(
    x = "Time (ms)",
    y = "Effect of Predictability (µV)",
    colour = "Picture-Word Congruency",
    fill = "Picture-Word Congruency"
  ) +
  theme(
    legend.key.height = unit(12, "pt"),
    legend.position = "top",
    axis.line.x = element_blank()
  )

ggsave(file.path("figs", "sample_level", "picture_word", "roi_simple_effs.pdf"), fe_se_pl, width=4.75, height=2.75, device="pdf")
ggsave(file.path("figs", "sample_level", "picture_word", "roi_simple_effs.png"), fe_se_pl, width=4.75, height=2.75, device="png", type="cairo")