# this script summarises the ERPs from the localiser task

library(lme4)
library(dplyr)
library(purrr)
library(readr)
library(tidyr)
library(ggplot2)
library(parallel)
library(patchwork)
library(svglite)

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

eff_cols <- c(
  "none" = "#000000",
  "Words vs. False Font" = "#D81B60",
  "Words vs. Phase-Shuffled" = "#1E88E5"
)

cond_cols <- c(
  "Words" = "#000000",
  "False Font" = "#D81B60",
  "Phase-Shuffled" = "#1E88E5"
)

# how many cores to parallelise over
n_cores <- 12

# get alist with the unique values for the times and channel names (taken from the first csv file)
unique_ids <- list.files("localiser_sample_data", pattern=".+\\.csv$", full.names=TRUE) %>%
  first() %>%
  read_csv() %>%
  select(ch_name, time) %>%
  as.list() %>%
  lapply(unique)

# function to import all data for selected channels at the selected times (exact matches)
get_sample_dat <- function(channels=NA, times=NA) {
  list.files("localiser_sample_data", pattern=".+\\.csv$", full.names=TRUE) %>%
    map_dfr(function(f) {
      read_csv(
        f,
        col_types=cols(
          subj_id = col_character(),
          stim_grp = col_integer(),
          resp_grp = col_integer(),
          trl_nr = col_integer(),
          ch_name = col_character(),
          condition = col_character(),
          string = col_character(),
          item_nr = col_integer(),
          .default = col_double()
        ),
        progress = FALSE
      ) %>%
        when(
          is.na(channels) ~ .,
          ~ filter(., ch_name %in% channels)
        ) %>%
        when(
          is.na(times) ~ .,
          ~ filter(., time %in% times)
        ) %>%
        # get a unique id for each stimulus and simplify condition column
        mutate(
          condition = substr(condition, 1, 1),
          item_id = paste(item_nr, condition, sep="_")
        ) %>%
        # remove unused columns to save memory
        select(-string, -resp_grp, -stim_grp, -trl_nr)
    }) %>%
    # deviation-code condition (with word as the null)
    mutate(
      ff_dev = scale(ifelse(condition=="b", 1, 0), center=TRUE, scale=FALSE),
      noise_dev = scale(ifelse(condition=="n", 1, 0), center=TRUE, scale=FALSE)
    )
}

rois <- list(
  c("TP7", "CP5", "P5", "P7", "P9", "PO3", "PO7", "O1"),
  c("TP8", "CP6", "P6", "P8", "P10", "PO4", "PO8", "O2")
)

# get fixed effect results from models for each timepoint, for the ROI
message("Importing data")

d_i_list <- get_sample_dat(channels = c(rois[[1]], rois[[2]])) %>%
  mutate(
    hemis = ifelse(ch_name %in% rois[[1]], "l", "r"),
    hemis_dev = scale(ifelse(hemis == "l", 0, 1), center=TRUE, scale=FALSE)
  ) |>
  group_split(time)

gc()

message(sprintf("Fitting %g models to timecourse in parallel", length(d_i_list)))

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

m_fe <- parLapply(cl, d_i_list, function(d_i) {
  m <- lme4::lmer(
    uV ~ (ff_dev + noise_dev) * hemis_dev +
      (1 | subj_id) +
      (1 | ch_name:subj_id) +
      (1 | item_nr) +
      (1 | item_id),
    REML = FALSE,
    control = lme4::lmerControl(optimizer="bobyqa"),
    data = d_i
  )
  
  m %>%
    summary() %>%
    with(coefficients) %>%
    as_tibble(rownames = "fe")
}) %>%
  reduce(bind_rows) %>%
  mutate(time = rep(unique_ids$time, each=6))

stopCluster(cl)


fe_res_tidy <- m_fe %>%
  mutate(
    stim_eff_lab = factor(case_when(
      grepl("ff_dev", fe, fixed=TRUE) ~ "Words vs. False Font",
      grepl("noise_dev", fe, fixed=TRUE) ~ "Words vs. Phase-Shuffled",
      TRUE ~ "none"
    ), levels = c("none", "Words vs. False Font", "Words vs. Phase-Shuffled")),
    main_or_int = factor(ifelse(
      grepl(":", fe, fixed=TRUE),
      "Interaction",
      "Main Effect"
    ), levels = c("Main Effect", "Interaction")),
    eff_lab = factor(case_when(
      grepl(":", fe, fixed=TRUE) ~ "Stimulus * Hemisphere",
      fe == "hemis_dev" ~ "Hemisphere",
      fe == "(Intercept)" ~ "Intercept",
      TRUE ~ "Stimulus"
    ), levels = c("Intercept", "Hemisphere", "Stimulus", "Stimulus * Hemisphere")),
    bound_lower = Estimate - 1.96 * `Std. Error`,
    bound_upper = Estimate + 1.96 * `Std. Error`
  )

pl_eff <- fe_res_tidy %>%
  ggplot(aes(time, Estimate, group=stim_eff_lab, colour=stim_eff_lab, fill=stim_eff_lab)) +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  geom_ribbon(aes(ymin=bound_lower, ymax=bound_upper), size=0.2, alpha=0.5) +
  # geom_line(size=0.5) +
  facet_wrap(vars(eff_lab), nrow=2) +
  labs(x = "Time (ms)", y = "Effect Estimate (µV)", fill="Fixed Effect (95% CI)", colour="Fixed Effect (95% CI)", tag="a") +
  scale_x_continuous(breaks = seq(-200, 1000, 200), expand=c(0, 0)) +
  scale_y_continuous(breaks = scales::pretty_breaks(n = 5)) +
  scale_fill_manual(values = eff_cols, breaks = c("Words vs. False Font", "Words vs. Phase-Shuffled")) +
  scale_colour_manual(values = eff_cols, breaks = c("Words vs. False Font", "Words vs. Phase-Shuffled")) +
  theme(
    legend.position = "top",
    legend.key.height = unit(6, "pt"),
    strip.background = element_blank(),
    axis.line.x = element_blank()
  )


preds <- fe_res_tidy %>%
  select(time, fe, Estimate) %>%
  pivot_wider(names_from = fe, values_from = Estimate) %>%
  expand_grid(
    condition = c("w", "b", "n"),
    hemis = c("l", "r")
  ) %>%
  left_join(
    d_i_list[[1]] %>%
      select(condition, hemis, hemis_dev, ff_dev, noise_dev) %>%
      distinct() %>%
      rename(
        ff_dev_val = ff_dev,
        noise_dev_val = noise_dev,
        hemis_dev_val = hemis_dev
      ),
    by = c("condition", "hemis")
  ) %>%
  mutate(
    pred_uV = `(Intercept)` +
      ff_dev_val * ff_dev +
      noise_dev_val * noise_dev +
      hemis_dev_val * hemis_dev +
      ff_dev_val * hemis_dev_val * `ff_dev:hemis_dev` +
      noise_dev_val * hemis_dev_val * `noise_dev:hemis_dev`,
    cond_lab = factor(recode(
      condition,
      w = "Words",
      b = "False Font",
      n = "Phase-Shuffled"
    ), levels = c("Words", "False Font", "Phase-Shuffled")),
    hemis_lab = factor(
      ifelse(hemis=="l", "Left Hemisphere", "Right Hemisphere"),
      levels=c("Left Hemisphere", "Right Hemisphere")
    )
  )

pl_preds <- preds %>%
  ggplot(aes(time, pred_uV, colour=cond_lab)) +
  geom_vline(xintercept=0, linetype="dashed") +
  geom_hline(yintercept=0, linetype="dashed") +
  geom_line() +
  scale_x_continuous(breaks = seq(-200, 1000, 200), expand=c(0, 0)) +
  scale_y_continuous(breaks = scales::pretty_breaks(n = 5)) +
  scale_colour_manual(values = cond_cols) +
  labs(
    x = "Time (ms)",
    y = "Amplitude (µV)",
    colour = "Stimulus ERP",
    tag = "b"
  ) +
  theme(
    legend.position = "top",
    plot.margin = margin(),
    strip.background = element_blank(),
    axis.line.x = element_blank()
  ) +
  facet_wrap(vars(hemis_lab))

# function to get a ggplot's legend as a ggplot object that patchwork will accept
get_legend <- function(pl) {
  tmp <- ggplot_gtable(ggplot_build(pl))
  leg <- which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
  legend <- tmp$grobs[[leg]]
  legend
}

eff_leg <- get_legend(pl_eff)
stim_leg <- get_legend(pl_preds)

pl_leg <- wrap_plots(list(eff_leg, stim_leg), ncol=1)

pl <- (
  (pl_leg) /
    (pl_eff + theme(legend.position="none")) /
    (pl_preds + theme(legend.position="none"))
  ) +
  plot_layout(heights = c(0.1, 0.05, 1, 0.4))

ggsave(file.path("figs", "sample_level", "localiser", "roi.svg"), pl, width=6.5, height=6.5)
ggsave(file.path("figs", "sample_level", "localiser", "roi.pdf"), pl, width=6.5, height=6.5)