# Readme ----------------------------------------------------------------------------------------------------------
# 1.  **Info**:
#   This script will process the fMRI signal from mat files in config specified folder 
#   calculating pearson correlation coefficient, renaming the matlab ROI with
#   Allan Brain atlas one. Then testing one-sample t-test for each edge within a given group.
#   This will also calculate the mean of correlation per group and detect outliers.
# 
# 2.  **Inputs**:
#   Config specification of folders/subfolders with mat files inside them
# 
# 3.  **Return**:
#   Data frames for:
#   + Rho correlation and fisher r-to-z transform for **upper diagonal matrix** on animal basis
#   + Rho correlation, fisher r-to-z transform and one-sample t-Test for **upper diagonal matrix** on group basis

# load libraries
# library(R.matlab)
# library(tidyverse)
# library(psych)
# library(multtest)
# library(Hmisc)

# load source files
# source("./R/data_functions.R")

# Process inputs --------------------------------------------------------------------------------------------------
# Calculate the correlation matrix for each animal individually, from mat files.
# Then combine all into a big table, from wish FisherZ will be calculated (R-to-Z Fisher tranform).
# Finally on a group level a one sample t-test will be calculated to get significant correlation per group.

process_inputs_raw <- function(raw.data.folder, 
                               Treatments.groups, Impulsivity.groups, 
                               Labels.data, Matlab.labels.file, 
                               permute = FALSE, 
                               remove.areas) {
  tryCatch({

# Load brain names data -------------------------------------------------------------------------------------------
    # Load the brain names data, matlab labels and heatmap colors for brain areas,
    # that was previously extracted from the matlab file from Peter.
    # load labels data to change the matlab annotation to Allan Brain atlas one later on.
    flog.info("Loading labels", name = log.name)
    Labels.data$fMRIName <- gsub(pattern = " ", replacement = "_", x = Labels.data$fMRIName)
    
    # Matlab brain labels
    Matlab.brain.labels <- scan(file = Matlab.labels.file, what = character()) 
    Matlab.brain.labels <- Matlab.brain.labels[-1]
    
# Run iterations --------------------------------------------------------------------------------------------------
    flog.info("Processing files", name = log.name)
    Data.raw <- data.frame()
    
    # Load files, calculate Pearson correlations and FisherZ score.
    for (i in seq_along(Treatments.groups)) {
      for (j in seq_along(Impulsivity.groups)) {
        
        file.list <- dir(path = paste(raw.data.folder, 
                                      Treatments.groups[i], 
                                      Impulsivity.groups[j], 
                                      sep = "/"), 
                         pattern = ".mat")
        
        flog.info(paste("Running:", Treatments.groups[i], Impulsivity.groups[j], sep = " "), name = log.name)
        
        # read mat file
        for (k in seq_along(file.list)) {
          
          Data.temp <- loadMatCorr(file = file.list[k], raw.data.folder = raw.data.folder,
                                   Treatment = Treatments.groups[i], 
                                   Impulsivity = Impulsivity.groups[j], 
                                   Matlab.brain.labels = Matlab.brain.labels ,
                                   remove.areas = remove.areas)  
          # gather
          Data.raw <- bind_rows(Data.raw, Data.temp)
          
        } # file list iteration
        
      } # impulvivity
    } # Treatments
    
    Data.temp <- NULL
    file.list <- NULL
    k <- NULL
    i <- NULL
    j <- NULL
    
    # Substitite the labels for short Allan brain atlas ones
    flog.info("Renaming labels", name = log.name)
    for (lab in 1:nrow(Labels.data)) {
      Data.raw$From <- gsub(Labels.data$fMRIName[lab], Labels.data$Atlas[lab], Data.raw$From)
      Data.raw$To <- gsub(Labels.data$fMRIName[lab], Labels.data$Atlas[lab], Data.raw$To)
    }
    
    # Do fisher Z transform 
    flog.debug("Fisher z tranform", name = log.name)
    Data.raw$FisherZ <- fisherz(Data.raw$correlation.rho)
    
    if (permute) {
      
      flog.info("Permuting", name = log.name)
      permute <- Data.raw %>% 
        select(Impulsivity, File) %>%
        unique(.) %>%
        mutate(Impulsivity = sample(Impulsivity))
      
      Data.raw <- Data.raw %>% 
        select(-Impulsivity) %>% 
        left_join(., permute, by = c("File"))
      
    }
    
    return(Data.raw)
    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}

# Run one sample --------------------------------------------------------------------------------------------------
# On a group level a one sample t-test will be calculated to get significant correlation per group.
# This will also calculate the mean of correlation and Fisher score per group

process_inputs_onesample <- function(Data.raw, PI_mu.set, PI_TSBH, PI_varEqual) {
  tryCatch({
    flog.info("Running one sample per each edge", name = log.name)
    
    Data.grouped <- Data.raw %>% 
      group_by(From, To, Impulsivity, Treatment) %>%
      summarise(OneSamplTtest.pval = t.test(x = FisherZ, 
                                            mu = PI_mu.set, 
                                            alternative = c("two.sided"), 
                                            var.eqaul = PI_varEqual)$p.value,
                OneSamplTtest.tval = as.numeric(t.test(x = FisherZ, 
                                                       mu = PI_mu.set, 
                                                       alternative = c("two.sided"), 
                                                       var.eqaul = PI_varEqual)$statistic),
                correlation.rho = mean(correlation.rho),
                FisherZ = mean(FisherZ))
    
    Data.grouped <- Data.grouped %>% 
      group_by(Impulsivity, Treatment) %>%
      mutate(Corrected.OneSamplTtest.pval = two.step.BH(OneSamplTtest.pval, method = PI_TSBH))
    
    return(Data.grouped)
    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}

# save files --------------------------------------------------------------------------------------------------

process_inputs_savefile <- function(Data.raw, Data.grouped, output.folder, PI_single.data, PI_grouped_data) {
  tryCatch({
    flog.info("Saving files", name = log.name)
    # Save data
    write_csv(x = Data.raw, 
              path = paste0(output.folder, PI_single.data, ".csv"))
    
    write_csv(x = Data.grouped, 
              path = paste0(output.folder, PI_grouped_data, ".csv"))
    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}