# Functions needed for plotting

# Heatmap plot ----------------------------------------------------------------------------------------------------
# Must have a full table to work on! Need to presort Heatmap.brains to contain only the needed areas.

plot_heatmap <- function(Data, measure = "adj.p.value", breaks = NA,
                         color = colorRampPalette(rev(brewer.pal(n = 7, name =
                                                                   "RdYlBu")))(100),
                         plot.name = "xxx",
                         Annotations.data = Heatmap.brains){

  # change the color of annotation to what you want: (eg: "navy", "darkgreen")
  Var1        <- Annotations.data$Color.simplyfied %>% unique(.)
  names(Var1) <- Annotations.data$Parental_brain_region %>% unique(.)
  Var2        <- c("navy", "darkgreen")
  names(Var2) <- unique(Annotations.data$Hemisphere)
  anno_colors <- list(Parental_brain_region = Var1, Hemisphere = Var2)
  
  annotation_rows <- Annotations.data %>% select(Parental_brain_region) %>% data.frame(.)
  rownames(annotation_rows) <- Annotations.data$Atlas
  
  annotation_columns <- Annotations.data %>% select(Hemisphere) %>% data.frame(.)
  rownames(annotation_columns) <- Annotations.data$Atlas
  
  # init matrix
  mat.size <- length(unique(Annotations.data$Atlas))
  
  mat <- matrix(NA, nrow = mat.size, ncol = mat.size) 
  
  dimnames(mat) = list(
    Annotations.data$Atlas, # row names 
    Annotations.data$Atlas) # col names
  
  # fill data
  
  mat[as.matrix(Data[c("From", "To")])] <- Data[[measure]]
  
  plot <- pheatmap(mat,
                   breaks = breaks,
                   color = color,
                   border_color = "grey60",
                   cluster_rows = FALSE, cluster_cols = FALSE,
                   annotation_col = annotation_columns,
                   annotation_row = annotation_rows,
                   annotation_colors = anno_colors,
                   annotation_names_row = FALSE,
                   annotation_names_col = FALSE,
                   silent = FALSE)[[4]]
  
  ggplot2::ggsave(file = paste0(plots.folder, plot.name, "_", measure, ".png"), 
                  plot = plot, 
                  type = "cairo-png",
                  width = 15,  height = 15, units = "in")
  
  ggplot2::ggsave(file = paste0(plots.folder, plot.name, "_", measure, ".pdf"), 
                  plot = plot,
                  width = 15,  height = 15, units = "in")
  
}


# Network graph ---------------------------------------------------------------------------------------------------

plot_network_gg <- function(Data, layoutMat, filter, plot.name = "Network_"){
  # This needs specific input like Linetype for the edges (positive, negative or no change)
  # color code for correlation as well as breaks for the correlations.
  
  Edge.data <- Data %>% 
    select(From, To, LineType, correlation.rho.col, correlation.rho.brk) %>%
    arrange(From)
  
  Node.data <- Data %>% 
    rename(From2 = From, To2 = To) %>%
    rename(From = To2, To = From2) %>%
    bind_rows(Data,.) %>%
    select(From, To, correlation.rho) %>%
    group_by(From) %>%
    mutate(Nodes.size = abs(mean(correlation.rho))) %>%
    select(From, Nodes.size) %>%
    unique(.) %>%
    arrange(From)
   
  Net.plot.data <- graph_from_data_frame(d = Edge.data, directed = FALSE, vertices = Node.data) %>% 
    ggnetwork(., layout = layoutMat)

  net.graph <- ggplot(data = Net.plot.data, 
                      aes(x = x, y = y, xend = xend, yend = yend)) +
    geom_edges(aes(color = correlation.rho.col, linetype = LineType),
               alpha = 0.5) +
    geom_nodelabel(aes(label = vertex.names, size = Nodes.size), fontface = "bold") + 
    scale_color_identity() +
    scale_linetype_manual(values = c("longdash", "blank", "solid")) +
    theme_blank() +
    theme(legend.position = "bottom")

  ggplot2::ggsave(filename = 
           paste0(plots.folder, plot.name, filter, ".png"), 
         plot = net.graph, 
         width = 15, height = 15, 
         units = "in", type = "cairo-png")
  
  ggplot2::ggsave(filename = 
                    paste0(plots.folder, plot.name, filter, ".pdf"), 
                  plot = net.graph, 
                  width = 15, height = 15, 
                  units = "in")
  
}

# Bar graph -------------------------------------------------------------------------------------------------------

Bar_graph <- function(Data, x.axis = "From", 
                      x.axis.title = "",
                      y.axis = "N",
                      y.axis.title = "Number of significantly different correlations",
                      color.fill = "N",
                      filename = "xxx", 
                      plot.title = "Change in correlation number for High vs Low", 
                      plots.folder = "./Plots/", 
                      output.folder = "./Output/"){
  if(nrow(Data) < 2){
    summ <- ggplot(data = data.frame()) + 
      geom_point() + xlim(0,10) + ylim(0,10) + 
      annotate("text", x = 5, y = 5, label = "Empty!")
  }
  
  summ <- ggplot(data = Data, 
                 aes_string(x = x.axis, y = y.axis, fill = color.fill)) + 
    geom_bar(stat = "identity", position = "dodge") + 
    scale_fill_distiller(palette = "RdYlBu") + 
    ggtitle(paste0(plot.title)) +
    xlab(x.axis.title) +
    ylab(y.axis.title) + 
    theme_bw() +
    theme(#legend.position = "none",
          #panel.grid.major = element_blank(),
          #panel.grid.minor = element_blank(),
          axis.text.x  = element_text(size = 18, angle = 90, hjust = 1, face = "bold", vjust = 0.4), 
          axis.text.y  = element_text(size = 16),
          axis.title.y = element_text(face = "bold", colour = "black", size = 22, vjust = 0.4),
          plot.title = element_text(hjust = 0, size = 22)
    )
  
  ggsave(filename = 
           paste0(plots.folder, filename, ".png"),
         plot = summ, 
         width = 15, height = 15, 
         units = "in", type = "cairo-png")
  
  ggsave(filename = 
           paste0(plots.folder, filename, ".pdf"),
         plot = summ, 
         width = 15, height = 15, 
         units = "in")
  
}
