# Readme ----------------------------------------------------------------------------------------------------------
# 1.  **Info**:
#   This script will plot:
#   + Heatmaps for each group as well as delta between the groups for correlation, FisherZ, tval for FisherZ
#   + Network graphs
#   + Violin plots for distributions
# 
# 2.  **Inputs**:
#   Config specification for output folders and file names. Single and grouped data frame from *Process inputs* script.
# 
# 3.  **Return**:
#   Heatmaps and brain networks for:
#   + Rho correlation and fisher r-to-z transform
#   + Rho correlation, fisher r-to-z transform and one-sample t-Test for on group basis

# load libraries

# library(pheatmap)
# library(ggnetwork)
# library(network)
# library(sna)
# library(igraph)
# library(intergraph)
# library(scales)
# library(RColorBrewer)
# library(tidyverse)
# library(broom)

# Grouped data heatmap -----------------------------------------------------------------------------------------------
# Produce grouped heatmaps for p values for corrected One sample t test, 
# t values for that as well as correlation coefficient.

GD_make_heatmap <- function(Data.grouped, Heatmap.brains.sorted, 
                            GD_heatmap, GD_delta_heatmap, GD.p.filter) {
  tryCatch({

# calculate the IQR --------------------------------------------------------    
    # Make the matrix full, needed for heatmap
    Data.heatmap <- Data.grouped  %>% 
      rename(From2 = From, To2 = To) %>%
      rename(From = To2, To = From2) %>%
      bind_rows(Data.grouped, .) %>% 
      unite(loop, Treatment, Impulsivity, sep = "_", remove = FALSE)
    
    IQR.correlation.rho <- Data.heatmap %>%
      pull(correlation.rho) %>% 
      quantile(., prob = Heatmap.color.range, type = 7) %>%
      round(., digits = 2)
    
    OneSamplTtest.tval <- Data.heatmap %>%
      pull(OneSamplTtest.tval) %>% 
      quantile(., prob = Heatmap.color.range, type = 7) %>%
      round(., digits = 2)
    
# Make grouped heatmap  --------------------------------------     
    for (filter in unique(Data.heatmap$loop)) {
      flog.debug(filter, name = log.name)
      
      Data.temp <- Data.heatmap %>% filter(loop == filter) %>% arrange(From, To)
      
      plot_heatmap(Data.temp,
                   measure = "Corrected.OneSamplTtest.pval",
                   breaks = c(0, GD.p.filter, 1),
                   color = c("red", "black"),
                   plot.name = paste(GD_heatmap, filter, 
                                     "OneSamplTtest.pval", sep = "_"),
                   Annotations.data = Heatmap.brains.sorted)
      
      plot_heatmap(Data.temp,
                   measure = "OneSamplTtest.tval",
                   breaks = OneSamplTtest.tval,
                   color = rev(brewer.pal(n = 11, name = "RdYlBu")),
                   plot.name = paste(GD_heatmap, filter, 
                                     "OneSamplTtest.tval", sep = "_"),
                   Annotations.data = Heatmap.brains.sorted)
      
      plot_heatmap(Data.temp,
                   measure = "correlation.rho",
                   breaks = IQR.correlation.rho,
                   color = rev(brewer.pal(n = 11, name = "RdYlBu")),
                   plot.name = paste(GD_heatmap, filter, 
                                     "correlation.rho", sep = "_"),
                   Annotations.data = Heatmap.brains.sorted)
      
    } # loop over    
    
# Make High vs low grouped heatmap  --------------------------------------     
# Delta high vs low
# Uses the same IQR as group-wise plots.
# Uses the same data frame as in heatmap. Note it should be a full matrix!
    flog.debug("Calculating delta high vs low", name = log.name)
    
    Delta.heatmap <- Data.heatmap %>%
      select(-loop) %>%
      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)
    
    flog.debug("Making delta heatmaps high vs low", name = log.name)
    
    for (filter in unique(Delta.heatmap$Treatment)) {
      flog.debug(filter, name = log.name)
      
      Data.temp <- Delta.heatmap %>% filter(Treatment == filter) %>% arrange(From, To)
      
      plot_heatmap(Data.temp,
                   measure = "correlation.rho",
                   breaks = IQR.correlation.rho,
                   color = rev(brewer.pal(n = 11, name = "RdYlBu")),
                   plot.name = paste(GD_delta_heatmap, filter, 
                                     "highVSlow",
                                     "correlation.rho", sep = "_"),
                   Annotations.data = Heatmap.brains.sorted)
      
    } # loop over
    
    
# Make Delta treatment grouped heatmap  --------------------------------------     
# Delta treatment Bi cmp vs Vehicle
# Uses the same IQR as group-wise plots.
# Uses the same data frame as in heatmap. Note it should be a full matrix!
    flog.debug("Calculating delta for treatment", name = log.name)
    
    BIvsVeh_Delta.heatmap <- Data.heatmap %>%
      select(-loop) %>%
      gather(key = "measure", value = "value", -From, -To, -Impulsivity, -Treatment) %>%
      spread(key = Treatment, value = value) %>%
      mutate(Delta = BI_compound - Vehicle) %>%
      select(-BI_compound, -Vehicle) %>%
      spread(key = measure, value = Delta)
    
    flog.debug("Making delta heatmaps for treatment", name = log.name)
    
    for (filter in unique(BIvsVeh_Delta.heatmap$Impulsivity)) {
      flog.debug(filter, name = log.name)
      
      Data.temp <- BIvsVeh_Delta.heatmap %>% filter(Impulsivity == filter) %>% arrange(From, To)
      
      plot_heatmap(Data.temp,
                   measure = "correlation.rho",
                   breaks = IQR.correlation.rho,
                   color = rev(brewer.pal(n = 11, name = "RdYlBu")),
                   plot.name = paste(GD_delta_heatmap, filter, 
                                     "BIcmpVSvehicle",
                                     "correlation.rho", sep = "_"),
                   Annotations.data = Heatmap.brains.sorted)
      
    } # loop over

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



# group_wise_network_graphs-----------------------------------------------------------------------------------------------
# Produce grouped heatmaps for p values for corrected One sample t test, 
# t values for that as well as correlation coefficient.

GD_make_network_graphs <- function(Data.grouped, Brain.coordinates.file, GD_network) {
  tryCatch({
    
    # Load rain coordinates --------------------------------------
    # Brain coordinates for network graph
    flog.debug("Loading brain coordinates", name = log.name)
    Brain.coordinates <- read_delim(file = Brain.coordinates.file, delim = ",")
    
    layoutMat <- Brain.coordinates %>% arrange(Node) %>%
      select(X,Y) %>% as.matrix(.)
    
    # Make line type positive or negative for the type of edge, also filter by corrected p.value
    # This does not require full matrix, as node data is precessed to full within the plot function
    flog.debug("Calculating network data", name = log.name)
    Data.network <- Data.grouped %>% 
      unite(loop, Treatment, Impulsivity, sep = "_", remove = FALSE) %>%
      mutate(LineType = case_when(.$OneSamplTtest.tval < 0 &
                                    .$Corrected.OneSamplTtest.pval < 0.05 ~ "Negative",
                                  .$OneSamplTtest.tval > 0 &
                                    .$Corrected.OneSamplTtest.pval < 0.05 ~ "Positive",
                                  TRUE ~ "No change"),
             # make binned correlation for edges (same category as for heatmaps)
             correlation.rho.brk = cut(correlation.rho, breaks = IQR.correlation.rho),
             correlation.rho.brk = gsub(pattern = "\\(", replacement = "", correlation.rho.brk),
             correlation.rho.brk = gsub(pattern = "\\]", replacement = "", correlation.rho.brk),
             correlation.rho.brk = gsub(pattern = ",", replacement = " to ", correlation.rho.brk),
             correlation.rho.col = cut(correlation.rho, 
                                       breaks = IQR.correlation.rho,
                                       labels = rev(brewer.pal(n = 11, name = "RdYlBu")))) %>%
      filter(!is.na(correlation.rho.col))
    
    flog.debug("Making network graphs", name = log.name)
    for (filter in unique(Data.network$loop)) {
      flog.debug(filter, name = log.name)
      Data.temp <- Data.network %>% filter(loop == filter)
      
      # use the same color range as in heatmap!
      net.graph <- plot_network_gg(Data.temp, layoutMat, filter = filter, plot.name = GD_network)
      
    } # loop over
    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}