# Readme ----------------------------------------------------------------------------------------------------------

# 1.  **Info**:
#   This script will plot:
#   + Calculate the unpaired t-test for each edge between high and low, report FDR corrected and uncorrected data (summary table plus box plot)
#   + Plot network graphs for each group (or correlation, filtered by significance)
#   + Calculate K.smirnov distribution test between high vs low for each brain area (also collapsed by hemisphere), corrected by FDR, make table and a plot with summary (bar graph with sum and violin plot with distribution)
#   + Plot network graphs for delta with significance filtering
# 
# 2.  **Inputs**:
#   Config specification for output folders and file names. Single and grouped data frame from *Process inputs* script.
# 
# 3.  **Return**:
#   CSV tables and appropriate box/violin plots:
#   + High vs low impulsive using fisher r-to-z transform on *each edge*. This will report heatmap of corrected and uncorrected p values,
#     as well as summary (total N) of significant edges for each brain regions (split and collapsed by hemisphere)
#   + Brain Network plot for high vs low (delta) with significance filtering (with and without correction)
#   + K.smirnov significance table for comparing distibutions between high and low; with boxplot/violin plot corrected by FDR.
#   + FDR corrected One sample t-Test for delta high-low on each node table; with boxplot/violin plot showing significant nodes.
# library(pheatmap)
# library(ggnetwork)
# library(network)
# library(sna)
# library(igraph)
# library(intergraph)
# library(scales)
# 
# library(RColorBrewer)
# 
# library(tidyverse)
# library(multtest)

# HvsL_edge_stats-----------------------------------------------------------------------------------------------
# Produce grouped heatmap

HvsL_edge_stats <- function(Data.single, Data.grouped,
                            Heatmap.brains.sorted,
                            HvsL_p.correction, HvsL_p.alpha, HvsL_heatmap.p.filter, 
                            HvsL_heatmap, 
                            HvsL_edgewise.columns, HvsL_single.hemisphere,
                            HvsL_mean_FisherZ_box, HvsL_mean_Regionwise_tTest,
                            HvsL_mean_Regionwise_KS, HvsL_edgeTtest_make.plots) {
  tryCatch({
    
    # High vs Low comparison --------------------------------------------------------    
    # This will use the animal-wise data: Data.single; to calculate t-test on each edge.
    # **NOTE**: It uses FisherZ column for this calculation.

    ## Edge-wise t-Test: *This ends in insignificant contrasts*
    flog.debug("Edge-wise t-Test", name = log.name)
    HvsL.stats <- Data.single %>%
      mutate(Impulsivity = factor(Impulsivity, levels = Impulsivity.groups)) %>%
      group_by(Treatment, From, To) %>% 
      nest() %>%
      mutate(tTest = map(data, HvsL.tTest)) %>%
      select(-data) %>%
      unnest()
    
    # correct P by two-step BH
    flog.debug("Correct P by two-step BH", name = log.name)
    HvsL.stats <- HvsL.stats %>%
      group_by(Treatment) %>%
      mutate(corrected.p.value = two.step.BH(p.value, 
                                             method = HvsL_p.correction, 
                                             alpha = HvsL_p.alpha)) 
    
    # make fill matrix
    HvsL.heatmap <- HvsL.stats %>%
      rename(From2 = From, To2 = To) %>%
      rename(From = To2, To = From2) %>%
      bind_rows(HvsL.stats, .)
    
    if (HvsL_edgeTtest_make.plots) {
      flog.debug("Making edge T.test plots", name = log.name)
      
      for (filter in unique(HvsL.heatmap$Treatment)) {
        flog.debug(filter, name = log.name)
        
        Data.temp <- HvsL.heatmap %>% filter(Treatment == filter) %>% arrange(From, To)
        
        plot_heatmap(Data.temp,
                     measure = "corrected.p.value",
                     breaks = c(0, HvsL_heatmap.p.filter, 1),
                     color = c("red", "black"),
                     plot.name = paste(HvsL_heatmap, filter, 
                                       "corrected.p.value", sep = "_"),
                     Annotations.data = Heatmap.brains.sorted)
        
        plot_heatmap(Data.temp,
                     measure = "p.value",
                     breaks = c(0, HvsL_heatmap.p.filter, 1),
                     color = c("red", "black"),
                     plot.name = paste(HvsL_heatmap, filter, 
                                       "raw.p.value", sep = "_"),
                     Annotations.data = Heatmap.brains.sorted)
        
      } # loop over
      
      # Box plot with counts of significant edges form a given nodes. Also collapsed by hemisphere.
      # **NOTE**: uses the *full* matrix from heatmap.
      # **NOTE#2**: the filenames are automately generated.
      # *This is insignificant*
      
      flog.debug("Making edge box plots", name = log.name)
      for (column in HvsL_edgewise.columns) {
        for (split in HvsL_single.hemisphere) {
          for (Treat in unique(HvsL.heatmap$Treatment)) {
            
            temp <- ifelse(split, "merged.hemi", "single.nodes")
            
            file.name <- paste("HvsL_edgeT.test", Treat, gsub(pattern = "<", "_", column), temp, sep = "_")
            flog.debug(file.name, name = log.name)
            
            Summ <- sum_data(HvsL.heatmap %>% filter(Treatment == Treat), filter = column, single.hemisphere = split)
            From.order <- Summ %>% group_by(From) %>% summarise(n = sum(N))
            
            Summ <- Summ %>% mutate(From = factor(From, levels = From.order$From[order(-From.order$n)]))
            
            Bar_graph(Summ, x.axis = "From",
                      y.axis = "N",
                      color.fill = "N",
                      filename = file.name, 
                      plot.title = "T.test significant correlations number for High vs Low", 
                      plots.folder = plots.folder, 
                      output.folder = output.folder)
            flog.debug("writing file", name = log.name)
            write.csv(x = Summ, file = paste0(output.folder, file.name, ".csv"), row.names = FALSE)
          }
        }
      }
      
      
    } # if make plots
    
# Node collapsed edge-wise t-test
    # Take the entire "column" of FisherZ scores per region and compare high vs low impulsivity using t.test:
    #  Output: HighVSLow_Regionwise_tTest_all.csv
    #  + Plot boxplot and violin plot distributions
    #  + Make the boxplot for FisherZ for a given brain region, data taken a as a "column" from FC heatmap
    #  + Calculate the high vs low t-test for a given brain region using the all, "column", FisherZ score for that brain region.
    #  *This also is insignificant*
    #  **NOTE**: uses the *diagonal* matrix from Data grouped, and makes first full matrix.
    #  **NOTE#2:** This is done for collapsed hemisphere, can be done also with split
    
    Data.grouped.full <- Data.grouped %>%
      rename(From2 = From, To2 = To) %>%
      rename(From = To2, To = From2) %>%
      bind_rows(Data.grouped, .) %>%
      mutate(From = gsub(pattern = " L| R", replacement = "", From)) 
    
    # plot distributions as box plot and violin plot
    # Make the boxplot/violin for FisherZ for a given brain region, data taken a as a "column" from FC heatmap
    
    if (HvsL_edgeTtest_make.plots) {
      flog.debug("Making distribution plots plots", name = log.name)
      
      for (treat in unique(Data.grouped.full$Treatment)) {
        # make boxplot
        p.box <- ggplot(data = Data.grouped.full %>% filter(Treatment == treat), 
                        aes(x = From, y = FisherZ, fill = Impulsivity)) + geom_boxplot() + ggtitle(treat) +
          theme_bw()
        
        ggplot2::ggsave(filename = paste0(plots.folder, HvsL_mean_FisherZ_box, treat, "boxplot.png"),
                        plot = p.box,
                        width = 30, height = 15, 
                        units = "in", type = "cairo-png")
        
        # make violin plot
        p.box <- ggplot(data = Data.grouped.full %>% filter(Treatment == treat), 
                        aes(x = From, y = FisherZ, fill = Impulsivity)) + geom_violin() + ggtitle(treat) +
          theme_bw()
        
        ggplot2::ggsave(filename = paste0(plots.folder, HvsL_mean_FisherZ_box, treat, "violin.png"),
                        plot = p.box,
                        width = 30, height = 15, 
                        units = "in", type = "cairo-png")
      }
      
    }
    
    # Calculate the high vs low t-test for a given brain region using the all, "column", FisherZ score for that brain region.
    flog.debug("Calculating Regionwise_tTest", name = log.name)
    Regionwise_tTest <- Data.grouped.full %>% group_by(From, Treatment) %>%
      nest() %>% 
      mutate(tTest = map(data, HvsL.tTest)) %>%
      select(-data) %>%
      unnest() %>%
      mutate(sign = if_else(p.value < 0.05, "red", "black"))
    
    write_csv(Regionwise_tTest, path = paste0(output.folder, HvsL_mean_Regionwise_tTest, ".csv"))
    
    if (HvsL_edgeTtest_make.plots) {
      
      p <- ggplot(data = Regionwise_tTest, aes(x = From, y = p.value, fill = sign)) + geom_bar(stat = "identity") +
        scale_fill_identity() +
        facet_grid(Treatment ~ .)
      
      ggplot2::ggsave(filename = paste0(plots.folder, HvsL_mean_Regionwise_tTest, ".png"),
                      plot = p,
                      width = 25, height = 15, 
                      units = "in", type = "cairo-png")
    }
    
    
    
    
    # Kolmogorov-Smirnov distribution comparison  --------------------------------------     
    # This is an another idea to compare the nodes.
    #  **Rationale:** We want to compare the FisherZ (correlation coefficient) for each node, 
    # and pick up the most different ones for high impulsive vs low impulsive animals. 
    # If the distributions differ in quality, and/or mode of distribution (i.e. uniform vs bimodal) 
    # K.S will pick that up, but t-test done above would not.
    flog.debug("Kolmogorov-Smirnov distribution", name = log.name)
    
    HvsL.KStest <- Data.grouped %>%
      rename(From2 = From, To2 = To) %>%
      rename(From = To2, To = From2) %>%
      bind_rows(Data.grouped, .) %>%
      mutate(From_both = gsub(pattern = " L| R", replacement = "", From)) %>% # can be taken out to have split between hemispheres
      group_by(Treatment, From_both) %>%
      nest() %>%
      mutate(KS.pvalue = map(data, KSmirnov.test)) %>% select(-data) %>%
      unnest() %>%
      group_by(Treatment) %>%
      mutate(KS.Corrected.pval = two.step.BH(KS.p.value, method = "TSBH")) %>%
      rename(From = From_both)
    
    # write data to file
    write_csv(x = HvsL.KStest, path = paste0(output.folder, HvsL_mean_Regionwise_KS, ".csv"))
    
    if (HvsL_edgeTtest_make.plots) {
      flog.debug("Making Kolmogorov-Smirnov distribution plots ", name = log.name)
      for (treat in unique(HvsL.KStest$Treatment)) {
        
        Temp <- HvsL.KStest %>% filter(Treatment == treat & KS.Corrected.pval < 0.05)
        
        From.order <- Temp %>% group_by(From) %>% summarise(n = mean(KS.D.value))
        
        Temp <- Temp %>% mutate(From = factor(From, levels = From.order$From[order(-From.order$n)]))
        
        Bar_graph(Temp, 
                  x.axis = "From",
                  y.axis.title = "Kolmogorov-Smirnov D value",
                  y.axis = "KS.D.value",
                  color.fill = "KS.Corrected.pval",
                  filename = paste0("HvsL_Kolmogorov.Smirnov_", treat), 
                  plot.title = "Kolmogorov-Smirnov for High vs Low, only significant shown", 
                  plots.folder = plots.folder, 
                  output.folder = output.folder)
      }
      
    } # if make plots
    
    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}


# Delta ----------------------------------------------------------------------------------------------
# This will use the group-wise mean data: **Data.grouped** to calculate delta for some measures, 
# and then perform a one sample t-test for each brain node, with FDR correction.

HvsL_delta_stats <- function(Data.grouped, HvsL_meanDelta_Regionwise_oneSample) {
  tryCatch({
    
    flog.debug("Calulating delta", name = log.name)
    
    Delta.grouped <- Data.grouped %>%
      gather(key = "measure", value = "value", -From, -To, -Impulsivity, -Treatment) %>%
      spread(key = Impulsivity, value = value) %>%
      mutate(Delta = high - low) %>%
      select(-high, -low) %>%
      spread(key = measure, value = Delta)
    
    # One sample t.test --------------------------------------------------------    
    # One sample t-test per hemisphere
    Delta.grouped.full <- Delta.grouped %>%
      rename(From2 = From, To2 = To) %>%
      rename(From = To2, To = From2) %>%
      bind_rows(Delta.grouped, .) %>%
      mutate(From = gsub(pattern = " L| R", replacement = "", From))
    
    # Calculate one sample t.Test
    flog.debug("Calulating delta one sample t.test", name = log.name)
    
    Delta.grouped.tTest <- Delta.grouped.full %>% 
      group_by(Treatment, From) %>% 
      summarise(OneSamplTtest.pval = t.test(x = FisherZ, 
                                            mu = 0, 
                                            alternative = c("two.sided"), 
                                            var.eqaul = TRUE)$p.value,
                OneSamplTtest.tval = as.numeric(t.test(x = FisherZ, 
                                                       mu = 0, 
                                                       alternative = c("two.sided"), 
                                                       var.eqaul = TRUE)$statistic),
                df = t.test(x = FisherZ, 
                            mu = 0, 
                            alternative = c("two.sided"), 
                            var.eqaul = TRUE)$parameter,
                Mean.delta.FisherZ = mean(FisherZ),
                Mean.delta.correlation = mean(correlation.rho)) %>%
      group_by(Treatment) %>%
      mutate(Corrected.OneSamplTtest.pval = two.step.BH(OneSamplTtest.pval, method = "TSBH"))
    
    # write data to file
    flog.debug("Saving delta one sample t.test", name = log.name)
    write_csv(x = Delta.grouped.tTest, path = paste0(output.folder, HvsL_meanDelta_Regionwise_oneSample, ".csv"))
    
    # summary plots  --------------------------------------
    if (HvsL_meanDelta_make.plots) {
      flog.debug("Making delta summary plots", name = log.name)
      
      for (treat in unique(Delta.grouped.tTest$Treatment)) {
        
        Temp <- Delta.grouped.tTest %>% filter(Treatment == treat & Corrected.OneSamplTtest.pval < 0.05)
        
        From.order <- Temp %>% group_by(From) %>% summarise(n = mean(OneSamplTtest.tval))
        
        Temp <- Temp %>% mutate(From = factor(From, levels = From.order$From[order(-From.order$n)]))
        
        Bar_graph(Temp, 
                  x.axis = "From",
                  y.axis = "OneSamplTtest.tval",
                  y.axis.title = "One Sample t-Test T value",
                  color.fill = "OneSamplTtest.tval",
                  filename = paste0("HvsL_Delta_OneSamplTtest", treat), 
                  plot.title = "One sample tTest for Delta for High vs Low, only significant shown", 
                  plots.folder = plots.folder, 
                  output.folder = output.folder)
      } # loop over plots
      
    } # if make plot
    
    return(Delta.grouped.tTest)

    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}

