# README --------------------------------------------------------------------------------------------------

#  1.  **Info**:
# This script will plot:
#   + Fig1: One sample t value for the node (region) wise delta.
#   + Fig2: 2-way RM ANOVA F score for Treatment:Impulsivity interaction
# 
# 2.  **Inputs**:
#   Config specification for output folders and file names. Single and grouped data frame from *Process inputs* script.
#   Also needs the *HvsL_meanDelta_Regionwise_oneSample* from HvsL_analysis script and *Interaction_node* from Interaction_analysis script.
# 
# 3.  **Return**:
#   Bar Plots:
#   + *Fig1*: One sample t value for the node (region) wise delta.
#   + *Fig2*: 2-way RM ANOVA F score for Treatment:Impulsivity interaction
# 
# # Initialize
#   Load configs and fixed data.


# library(pheatmap)
# library(ggnetwork)
# library(network)
# library(sna)
# library(igraph)
# library(intergraph)
# library(scales)
# 
# library(RColorBrewer)
# library(tidyverse)
# 
# # load source files
# source("./R/data_functions.R")
# source("./R/plot_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).


make_paper_figures <- function(Data.single, Data.grouped, 
                               Interaction.node, Delta.OneSample,
                               Impulsivity.groups, Treatments.groups,
                               Heatmap.brains.sorted,
                               make_paper_p.adjust.method = "bonferroni") {
  tryCatch({
    flog.debug("Loading data", name = log.name)
    
    # Read single animal data and factor Impulsivity and Treatment
    Data.single <- Data.single %>% ungroup() %>%
      mutate(Impulsivity = factor(Impulsivity, levels = Impulsivity.groups),
             Treatment = factor(Treatment, levels = Treatments.groups))
    
    # Read single animal data and factor Impulsivity and Treatment
    Data.grouped <- Data.grouped %>% ungroup() %>% 
      mutate(Impulsivity = factor(Impulsivity, levels = Impulsivity.groups),
             Treatment = factor(Treatment, levels = Treatments.groups))

# Fig 1 -----------------------------------------------------------------------------------------------------------
    # Make figure 1. which will show results from One sample t Test for each node where the underaying data were:
    #   + Delta of FisherZ for high impulsive - low impulsive
    
    Delta.OneSample <- Delta.OneSample %>% ungroup() %>%
      mutate(Corrected.OneSamplTtest.pval = p.adjust(OneSamplTtest.pval, method = make_paper_p.adjust.method)) %>%
      mutate(Significant = if_else(Corrected.OneSamplTtest.pval < 0.05, "Significant", "Not_significant"))
    
    for (Treat in unique(Delta.OneSample$Treatment)) {
      
      Temp <- Delta.OneSample %>% filter(Treatment == Treat)
      
      Temp <- Temp %>% mutate(From = factor(From, levels = Temp$From[order(-Temp$OneSamplTtest.tval)])) %>%
        #group_by(Significant) %>% 
        arrange(From) %>%
        mutate(N = n(), Color = colorRampPalette(brewer.pal(11, "RdYlBu"))(N),
               Color = if_else(Significant == "Not_significant", "grey60", Color))
      
      
      plot <- ggplot(data = Temp, 
                     aes(x = From, y = OneSamplTtest.tval, fill = Color)) + 
        geom_bar(stat = "identity", position = "dodge") + 
        scale_fill_identity() +
        xlab("") +
        ylab("One sample T test: t value") + 
        theme_bw() +
        theme(panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              panel.background = element_blank(),
              panel.border = element_blank(),
              axis.line = element_line(colour = "black", size = 1),
              axis.text.x  = element_text(size = 28, angle = 90, hjust = 1, face = "bold", vjust = 0.4), 
              axis.text.y  = element_text(size = 24, face = "bold", vjust = 0.4),
              axis.title.y = element_text(face = "bold", colour = "black", size = 32, vjust = 0.4),
              axis.ticks = element_line(colour = "black", size = 1),
              plot.title = element_text(hjust = 0, size = 22))
      
      ggsave(filename = 
               paste0(plots.folder, "Figure1_OneSampleDelta_wide_bonf_", Treat, ".png"),
             plot = plot, 
             width = 25, height = 10, 
             units = "in", type = "cairo-png")
      
      ggsave(filename = 
               paste0(plots.folder, "Figure1_OneSampleDelta_wide_bonf_", Treat, ".pdf"),
             plot = plot, 
             width = 25, height = 10, 
             units = "in")
      
    }    
    
    # Supplementary Figure 1 -----------------------------------------------------------------------------------------------------
    # This will plot the dynamite plot (with/without single values) for the Figure 1
    # using delta of high - low for given treatment.  
    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)
    
    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)) %>%
      select(From, Treatment, To, correlation.rho, FisherZ)
    
    Delta.grouped.full <- left_join(Delta.grouped.full, Delta.OneSample, by = c("From", "Treatment")) 
    
    for (Treat in unique(Delta.OneSample$Treatment)){
      
      Temp <- Delta.OneSample %>% filter(Treatment == Treat) %>% 
        mutate(From = factor(From, levels = Temp$From[order(-Temp$OneSamplTtest.tval)])) %>%
        arrange(From) %>%
        mutate(N = n(), Color = colorRampPalette(brewer.pal(11, "RdYlBu"))(N),
               Color = if_else(Significant == "Not_significant", "grey60", Color))
      
      
      Temp.grouped.mean <- Delta.grouped.full %>% filter(Treatment == Treat) %>%
        group_by(From, Treatment) %>%
        summarise(MeanCor = mean(correlation.rho), SDCor = sd(correlation.rho), NCor = n(), SEMCor = SDCor/NCor,
                  y0 = quantile(correlation.rho, 0.10), y25 = quantile(correlation.rho, 0.25), y50 = median(correlation.rho), 
                  y75 = quantile(correlation.rho, 0.75), y100 = quantile(correlation.rho, 0.90),
                  MeanFish = mean(FisherZ), SDFish = sd(FisherZ), NFish = n(), SEMFish = SDFish/NFish,
                  y0f = quantile(FisherZ, 0.10), y25f = quantile(FisherZ, 0.25), y50f = median(FisherZ), 
                  y75f = quantile(FisherZ, 0.75), y100f = quantile(FisherZ, 0.90)) %>%
        ungroup() %>%
        mutate(Treatment = factor(Treatment, levels = rev(Treatments.groups)))
      
      Temp.grouped.mean <- left_join(Temp.grouped.mean, Temp, by = c("From", "Treatment")) %>% 
        mutate(From = factor(From, levels = Temp$From[order(-Temp$OneSamplTtest.tval)])) %>%
        arrange(From) %>%
        mutate(N = n(), Color = colorRampPalette(brewer.pal(11, "RdYlBu"))(N),
               Color = if_else(Significant == "Not_significant", "grey60", Color)) %>%
        mutate(p.text = case_when(
          Corrected.OneSamplTtest.pval >= 0.05 ~ "ns",
          Corrected.OneSamplTtest.pval < 0.05 & Corrected.OneSamplTtest.pval >= 0.01 ~ "*",
          Corrected.OneSamplTtest.pval < 0.01 & Corrected.OneSamplTtest.pval >= 0.001 ~ "**",
          Corrected.OneSamplTtest.pval < 0.001 & Corrected.OneSamplTtest.pval >= 0.0001 ~ "***",
          Corrected.OneSamplTtest.pval < 0.0001 ~ "****"
        ))
      
      plot <- ggplot(data = Temp.grouped.mean, 
                     aes(x = From, fill = Color)) + 
        geom_boxplot(aes(ymin = y0, lower = y25, middle = y50, 
                         upper = y75, ymax = y100), stat = "identity") + 
        geom_point(aes(y = MeanCor), color = "black", size = 3) +
        geom_text(aes(y = y100 * 1.05, label = p.text), fontface = "bold", size = 5) +
        geom_hline(yintercept = 0, linetype = "dashed", color = "black") +
        scale_fill_identity() +
        xlab("") +
        ylab("Correlation difference: high - low") + 
        theme_bw() +
        theme(panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              panel.background = element_blank(),
              panel.border = element_blank(),
              axis.line = element_line(colour = "black", size = 1),
              axis.text.x  = element_text(size = 28, angle = 90, hjust = 1, face = "bold", vjust = 0.4), 
              axis.text.y  = element_text(size = 24, face = "bold", vjust = 0.4),
              axis.title.y = element_text(face = "bold", colour = "black", size = 32, vjust = 0.4),
              axis.ticks = element_line(colour = "black", size = 1),
              plot.title = element_text(hjust = 0, size = 22))
      
      # sava

      ggsave(filename = 
               paste0(plots.folder, "Suppl.Fig1_", Treat, ".png"),
             plot = plot, 
             width = 25, height = 10, 
             units = "in", type = "cairo-png")
      
      ggsave(filename = 
               paste0(plots.folder, "Suppl.Fig1_", Treat, ".pdf"),
             plot = plot, 
             width = 25, height = 10, 
             units = "in")
      
    }
    
    # Fig2. ANOVA interaction ---------------------------------------------------------------------------------------
    # This will plot a similar plot, but for the ANOVA interaction  
    
    Interaction.node <- Interaction.node %>%
      group_by(Stat) %>%
      mutate(corrected.p.value = p.adjust(p.value, method = make_paper_p.adjust.method)) %>%
      ungroup() %>%
      mutate(Significant = if_else(corrected.p.value < 0.05, "Significant", "Not_significant")) %>%
      mutate(From = From_both)
    
    for(Treat in unique(Interaction.node$Stat)){
      
      Temp <- Interaction.node %>% filter(Stat == Treat)
      
      Temp <- Temp %>% mutate(From = factor(From, levels = Temp$From[order(-Temp$`F value`)])) %>%
        arrange(From) %>%
        mutate(N = n(), Color = colorRampPalette(brewer.pal(11, "RdYlBu"))(N),
               Color = if_else(Significant == "Not_significant", "grey60", Color))
      
      
      plot <- ggplot(data = Temp, 
                     aes(x = From, y = `F value`, fill = Color)) + 
        geom_bar(stat = "identity", position = "dodge") + 
        scale_fill_identity() +
        xlab("") +
        ylab("2-way RM ANOVA: F value") +
        theme_bw() +
        theme(#legend.position = "none",
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          panel.background = element_blank(),
          panel.border = element_blank(),
          axis.line = element_line(colour = "black", size = 1),
          axis.text.x  = element_text(size = 28, angle = 90, hjust = 1, face = "bold", vjust = 0.4), 
          axis.text.y  = element_text(size = 24, face = "bold", vjust = 0.4),
          axis.title.y = element_text(face = "bold", colour = "black", size = 32, vjust = 0.4),
          axis.ticks = element_line(colour = "black", size = 1),
          plot.title = element_text(hjust = 0, size = 22))
      
      ggsave(filename = 
               paste0(plots.folder, "Figure2_2wayRM_ANOVA_wide_bonf", Treat, ".png"),
             plot = plot, 
             width = 25, height = 10, 
             units = "in", type = "cairo-png")
      
      ggsave(filename = 
               paste0(plots.folder, "Figure2_2wayRM_ANOVA_wide_bonf", Treat, ".pdf"),
             plot = plot, 
             width = 25, height = 10, 
             units = "in")
    }
    
    
    # Suppl Figure 2 ---------------------------------------------------------------------------------------
    # This will plot the supplementary figure 2, so raw correlation values for STN as node 
    # with rest of the brain network for high/low x Bicmp11/veh.
    # In total will be 4 bars/violin plots.  
    
    
    Node.full <-  Data.single %>%
      rename(From2 = From, To2 = To) %>%
      rename(From = To2, To = From2) %>%
      bind_rows(Data.single, .) %>%
      mutate(File = factor(paste(From, To, File, sep = "_"))) %>%
      mutate(From = gsub(pattern = " L| R", replacement = "", From)) %>%
      select(From, To, Treatment, Impulsivity, correlation.rho, FisherZ, File)
    
    #  model <- aov(FisherZ ~ Impulsivity * Treatment + Error(File/Treatment), data = STN.full)
    Node.order <- Interaction.node %>% 
      filter(Stat == "Impulsivity_Treatment") %>% 
      arrange(desc(`F value`)) %>% #sort largest/smallest
      pull(From) 
    
    Node.mean <- Node.full %>%
      group_by(From, Treatment, Impulsivity) %>%
      summarise(MeanCor = mean(correlation.rho), SDCor = sd(correlation.rho), NCor = n(), SEMCor = SDCor/NCor,
                y0 = quantile(correlation.rho, 0.10), y25 = quantile(correlation.rho, 0.25), y50 = median(correlation.rho), 
                y75 = quantile(correlation.rho, 0.75), y100 = quantile(correlation.rho, 0.90),
                MeanFish = mean(FisherZ), SDFish = sd(FisherZ), NFish = n(), SEMFish = SDFish/NFish,
                y0f = quantile(FisherZ, 0.10), y25f = quantile(FisherZ, 0.25), y50f = median(FisherZ), 
                y75f = quantile(FisherZ, 0.75), y100f = quantile(FisherZ, 0.90)) %>%
      ungroup() %>%
      mutate(Treatment = factor(Treatment, levels = rev(Treatments.groups)),
             Impulsivity = factor(Impulsivity, levels = rev(Impulsivity.groups)),
             From = factor(From, levels = Node.order))
    
    plotCor <- ggplot(data = Node.mean, aes(x = From,
                                            fill = Treatment)) + 
      geom_boxplot(aes(ymin = y0, lower = y25, middle = y50, upper = y75, ymax = y100),
                   stat = "identity") + 
      geom_point(aes(y = MeanCor), color = "black", position = position_dodge(width = 0.9)) +
      scale_fill_manual(values = c("#D3D3D3", "#FDAE61")) +
      theme_bw() + 
      ylab("Correlation") +
      xlab("") +
      theme(#legend.position = "none",
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_blank(),
        axis.line = element_line(colour = "black", size = 1),
        axis.text.x  = element_text(size = 28, face = "bold", colour = "black", angle = 90, hjust = 1, vjust = 0.4), 
        axis.text.y  = element_text(size = 24, face = "bold", vjust = 0.4, colour = "black"),
        axis.title.y = element_text(face = "bold", colour = "black", size = 32, vjust = 0.4),
        axis.ticks = element_line(colour = "black", size = 1),
        plot.title = element_text(hjust = 0, size = 22),
        strip.background = element_blank(),
        strip.text.y = element_text(size = 28, colour = "black", face = "bold")) +
      facet_grid(Impulsivity~.)
    
    ggsave(filename = 
             paste0(plots.folder, "Suppl.Figure2_HvsL_Treatment", ".png"),
           plot = plotCor, 
           width = 25, height = 7, 
           units = "in", type = "cairo-png")
    
    ggsave(filename = 
             paste0(plots.folder, "Suppl.Figure2_HvsL_Treatment", ".pdf"),
           plot = plotCor, 
           width = 25, height = 7, 
           units = "in")
    
    
  }, error = function(e) {
    flog.error("ERROR!", name = log.name)
    flog.error(e, name = log.name)
  })
}
