#' Internal function for colors 
#'
#' This function return a color palette with the number of colors specified by n
#'
#' @param n number of colors needed
#' @return a vector with colors
my_palette <- function(n) {
  palette <- c("#3B4992", "#DF8F44", "#B24745", "#1B5E20", "#6B452B", "#8F7700",
               "#808180", "#91D1C2","#FABFD2")
  return(palette[seq_len(n)])
}

#' Plot differential regressions for a link
#' 
#' Plot differential regressions for any target-target pair in an omic dataset
#'
#' @param deggs_object an object of class `deggs` generated by
#' `get_diffNetworks`
#' @param assayDataName name of the assayData of interest. If an unnamed list of 
#' data was given to `get_diffNetworks`, the assayDataName here will be the 
#' number indicating the position of the data in the assayDataList provided
#' before (i.e. if the user wants to plot a differential interaction observed in 
#' the transcriptomic data, which was second in the list, then assayDataName 
#' must be 2, if only one data table was provided assayDataName must be 1). 
#' Default 1.
#' @param gene_A character. Name of the first target (gene, protein, metabolite,
#' etc.)
#' @param gene_B character. Name of the second target (gene, protein, metabolite,
#' etc.)
#' @param title plot title. If NULL (default), the name of the assayData will be 
#' used. Use empty character "" for no title. 
#' @param legend_position position of the legend in the plot. It can be
#' specified by keyword or in any parameter accepted by `xy.coords` (defalut
#' "topright")
#' @importFrom methods is
#' @importFrom grDevices adjustcolor
#' @importFrom graphics abline axis boxplot legend mtext points polygon text
#' @return base graphics plot showing differential regressions across 
#' categories. The p value of the interaction term of
#' gene A ~ gene B \* category is reported on top.
#' @examples
#' data("synthetic_metadata")
#' data("synthetic_rnaseqData")
#' data("synthetic_proteomicData")
#' data("synthetic_OlinkData")
#' assayData_list <- list("RNAseq" = synthetic_rnaseqData,
#'                        "Proteomics" = synthetic_proteomicData,
#'                        "Olink" = synthetic_OlinkData)
#' deggs_object <- get_diffNetworks(assayData = assayData_list,
#'                                  metadata = synthetic_metadata,
#'                                  category_variable = "response",
#'                                  regression_method = "lm",
#'                                  padj_method = "bonferroni",
#'                                  verbose = FALSE,
#'                                  show_progressBar = FALSE,
#'                                  cores = 1)
#' plot_regressions(deggs_object,
#'                  assayDataName = "RNAseq",
#'                  gene_A = "MTOR", 
#'                  gene_B = "AKT2",
#'                  legend_position = "bottomright")
#' @export
plot_regressions <- function(deggs_object,
                             assayDataName = 1,
                             gene_A,
                             gene_B,
                             title = NULL, 
                             legend_position = "topright") {
  
  if (!is(deggs_object, "deggs")) {
    stop("deggs_object must be of class deggs")
  }
  
  if (!(is.character(gene_A) && is.character(gene_B))) {
    stop("Both GeneA and GeneB must be character.")
  }
  
  sig_var <- ifelse(deggs_object[["padj_method"]] == "none", "p.value", "p.adj")
  metadata <- deggs_object[["metadata"]]
  assayData <- deggs_object[["assayData"]][[assayDataName]]
  categories <- deggs_object[["category_subset"]]
  regression_method <- deggs_object[["regression_method"]]
  
  if (!gene_A %in% rownames(assayData)) (
    stop("gene_A is not in rownames(", assayDataName, ")")
  )
  
  if (!gene_B %in% rownames(assayData)) (
    stop("gene_B is not in rownames(", assayDataName, ")")
  )
  
  metadata <- metadata[colnames(assayData)]
  
  if(length(unique(metadata)) == 1) (
    stop("All sample IDs in ", assayDataName, " belong to one category. 
         No differential analysis is possible.")
  )
  
  if (is.null(categories)) {
    categories <- levels(metadata)
  }
  
  category_length <- length(categories)
  if(is.null(title)) (
    title <- ifelse(is.numeric(assayDataName), 
                    names(deggs_object[["diffNetworks"]])[assayDataName],
                    assayDataName)
  )
  
  # prepare data frame
  # using both t() and as.vector() to be compatible with both matrices and dfs
  df <- data.frame(as.vector(t(assayData[gene_A, ])),
                   as.vector(t(assayData[gene_B, ])),
                   metadata,
                   check.names = FALSE)
  colnames(df) <- c(gene_A, gene_B, "category")
  
  # compute gene-gene regression
  if (category_length == 2) {
    if (regression_method == "lm") {
      lmfit <- stats::lm(df[, 2] ~ df[, 1] * df[, 3])
      # i.e.: gene_A ~ gene_B * category
      p_interaction <- stats::coef(summary(lmfit))[4, 4]
      fit <- lapply(categories, function(i) {
        x <- df[df[, "category"] == i, 1]
        y <- df[df[, "category"] == i, 2]
        if (length(x) > 0 || length(y) > 0) return(stats::lm(y ~ x))
      })
    }
    if (regression_method == "rlm") {
      robustfit <- MASS::rlm(df[, 2] ~ df[, 1] * df[, 3])
      # i.e.: gene_A ~ gene_B * category
      p_interaction <- sfsmisc::f.robftest(robustfit, var = 3)$p.value
      fit <- lapply(categories, function(i) {
        x <- df[df[, "category"] == i, 1]
        y <- df[df[, "category"] == i, 2]
        if (length(x) > 0 || length(y) > 0) return(MASS::rlm(y ~ x))
      })
    }
  }
  if (category_length >= 3) {
    # one-way ANOVA
    # i.e.: gene_A ~ gene_B * category
    res_aov <- stats::aov(df[, 2] ~ df[, 1] * df[, 3], data = df)
    p_interaction <- summary(res_aov)[[1]][["Pr(>F)"]][3]
    fit <- lapply(categories, function(i) {
      x <- df[df[, "category"] == i, 1]
      y <- df[df[, "category"] == i, 2]
      if (regression_method == "lm") {
        if (length(x) > 0 || length(y) > 0) return(stats::lm(y ~ x));
      }
      if (regression_method == "rlm") {
        if (length(x) > 0 || length(y) > 0) return(MASS::rlm(y ~ x))
      }
    })
  }
  
  # Plot
  prefix <- ifelse(deggs_object[["padj_method"]] == "none", "P", "Padj")
  col <- my_palette(n = category_length)
  x_adj <- (max(df[, 1], na.rm = TRUE) - min(df[, 1], na.rm = TRUE)) * 0.05
  new_x <- seq(min(df[, 1], na.rm = TRUE) - x_adj,
               max(df[, 1], na.rm = TRUE) + x_adj,
               length.out = 100
  )
  
  # prediction of the fitted model
  preds <- lapply(fit, 
                  function(i) stats::predict(i,
                                             newdata = data.frame(x = new_x),
                                             interval = 'confidence')
                  )
  
  # we have the interaction p value for a single pair,
  # adjusted p values can be found in the deggs_object if padj_method was NOT
  # set to none
  if (deggs_object[["padj_method"]] != "none") {
    all_interactions <- do.call(rbind, 
                                deggs_object[["diffNetworks"]][[assayDataName]][categories])
    pair_index <- which(all_interactions$from == gene_A & 
                          all_interactions$to == gene_B)
    # if you don't find gene_A-gene_B, try gene_B-gene_A:
    if(length(pair_index) == 0)(
      pair_index <- which(all_interactions$from == gene_B & 
                            all_interactions$to == gene_A)
    )
    sig_interaction <- all_interactions[pair_index, sig_var]
  } else {
    sig_interaction <- p_interaction
  }
  
  plot(df[, 1], df[, 2],
       type = 'n', bty = 'l', las = 1, cex.axis = 1.1,
       font.main = 1, cex.lab = 1.3, xlab = colnames(df)[1],
       ylab = colnames(df)[2],
       main = title
  )
  
  for (i in seq_along(categories)) {
    # plot confidence intervals
    polygon(c(rev(new_x), new_x), c(rev(preds[[i]][, 3]), preds[[i]][, 2]),
            col = adjustcolor(col[i], alpha.f = 0.15), border = NA
    )
    
    # plot regression lines
    abline(fit[[i]], col = col[i], lwd = 1.5)
    
    # plot dots 
    cols <- col[df[df[, "category"] == categories[i], 3]]
    pch <- c(16:(16 + category_length - 1))[df[df[, "category"] ==
                                                 categories[i], 3]]
    row <- points(df[df[, "category"] == categories[i], 1], # x coordinates from gene_A 
                  df[df[, "category"] == categories[i], 2], # y coordinates from gene_B 
                  cex = 1.5,
                  pch = pch, col = adjustcolor(cols, alpha.f = 0.7)
    )
  }
  mtext(
    bquote(paste(
      .(prefix)["interaction"] * "=",
      .(format(sig_interaction, digits = 2))
    )),
    cex = 1.2, side = 3, adj = 0.04
  )
  legend(
    x = legend_position, legend = categories,
    col = col, lty = 1,
    bty = "o", box.lty = 0,
    cex = 0.8
  )
}

#' Boxplots of single nodes (genes,proteins, etc.)
#'
#' This function is for internal use of `View_diffnetworks`
#'
#' @param gene gene name  (must be in `rownames(assayData)`)
#' @param assayDataName name of the assayData of interest. If an unnamed list of 
#' data was given to `get_diffNetworks`, the assayDataName here will be the 
#' number indicating the position of the data in the assayDataList provided
#' before (i.e. if the user wants to plot a differential interaction observed in 
#' the transcriptomic data, which was second in the list, then assayDataName 
#' must be 2, if only one data table was provided assayDataName must be 1). 
#' Default 1.
#' @param deggs_object  an object of class `deggs` generated by 
#' `get_diffNetworks`
#' @return the boxplot
node_boxplot <- function(gene,
                         assayDataName = 1,
                         deggs_object) {
  
  title <- ifelse(is.numeric(assayDataName), 
                  names(deggs_object[["diffNetworks"]])[assayDataName],
                  assayDataName)
  
  # if the selected node is not in present in the first assayData matrix 
  # show an empty plot with an informative message
  if (!gene %in% rownames(deggs_object[["assayData"]][[assayDataName]])) {
    plot(x = 0:10, y = 0:10, ann = FALSE,bty = "n",type = "n",
         xaxt = "n", yaxt = "n")
    text(x = 5, y = 9, labels = paste(gene, "is not present \n in", title), 
         font = 2, cex = 1.3)
    return()
  }
  metadata <- deggs_object[["metadata"]]
  x <- metadata[colnames(deggs_object[["assayData"]][[assayDataName]])]
  y <- as.numeric(deggs_object[["assayData"]][[assayDataName]][gene, ])
  col <- my_palette(n = nlevels(metadata))
  cols <- col[as.numeric(x)]
  
  boxplot(y ~ x,
          outline = FALSE, whisklty = 1, medlwd = 2, cex.axis = 1.1,
          col = NA, cex.lab = 1.3, ylab = gene, las = 2, boxwex = .5,
          ylim = range(y), xaxt = 'n', xlab = '',
          main = title)
  
  points(jitter(as.numeric(x), amount = 0.15), y,
         pch = 20, col = adjustcolor(cols, alpha.f = 0.7), cex = 2)
  
  xtick <- levels(x)
  axis(1, at = seq_along(xtick), labels = xtick, cex.axis = 1.2)
}


#' Interactive visualisation of differential networks
#'
#' Explore differential networks and interactively select regression 
#' and box plots
#'
#' @param deggs_object an object of class `deggs` generated by
#' `get_diffNetworks`
#' @param legend.arrow.width width of the arrow used in the network legend. 
#' Default is 0.35. As the number of assayData matrices increases this parameter
#' must be accordingly increased to avoid graphical errors in the legend. 
#' @param stepY_legend vertical space between legend arrows. It is used together 
#' with `legend.arrow.width` to adjust the legend space in case of graphical
#' errors. Default is 55. 
#' @import knitr
#' @import rmarkdown
#' @importFrom magrittr %>%
#' @return a shiny interface showing networks with selectable nodes and links
#' @examples 
#' data("synthetic_metadata")
#' data("synthetic_rnaseqData")
#' data("synthetic_proteomicData")
#' data("synthetic_OlinkData")
#' assayData_list <- list("RNAseq" = synthetic_rnaseqData,
#'                        "Proteomics" = synthetic_proteomicData,
#'                        "Olink" = synthetic_OlinkData)
#' deggs_object <- get_diffNetworks(assayData = assayData_list,
#'                                  metadata = synthetic_metadata,
#'                                  category_variable = "response",
#'                                  regression_method = "lm",
#'                                  verbose = FALSE,
#'                                  show_progressBar = FALSE,
#'                                  cores = 1)
#' # the below function runs a shiny app, so can't be run during R CMD check                                  
#' if(interactive()){
#' View_diffNetworks(deggs_object)  
#' }
#' @export
View_diffNetworks <- function(deggs_object,
                              legend.arrow.width = 0.35,
                              stepY_legend = 55) {
  
  if (!is(deggs_object, "deggs")) stop("deggs_object must be of class deggs")
  
  sig_var <- ifelse(deggs_object[["padj_method"]] == "none", "p.value", "p.adj")
  multiOmic <- ifelse(length(deggs_object[["assayData"]]) > 1, TRUE, FALSE)
  
  if (multiOmic) {
    category_networks <- get_multiOmics_diffNetworks(deggs_object)
  } else {
    category_networks <- deggs_object[["diffNetworks"]][[1]]
  }
  
  ################# server ########################
  server <- function(input, output, session) {
    empty_df <- data.frame("from" = NA, "to" = NA, "p.value" = 1, "p.adj" = 1, 
                           "layer" = NA)
    
    # check if the category_network has data 
    outVar <- shiny::reactive({
      if (!is.data.frame(category_networks[[input$category]])) return(empty_df)
      if (nrow(category_networks[[input$category]]) == 0) return(empty_df)
      return(category_networks[[input$category]])
    })
    
    nodes_selection <- shiny::reactiveValues(current_node = NULL)
    
    # Set up reactive table
    edges <- shiny::reactive({
      if (is.data.frame(category_networks[[input$category]])) {
        if (nrow(category_networks[[input$category]]) != 0) {
          edges <- category_networks[[input$category]]
          edges <- edges[edges[, sig_var] < input$slider, ]
          edges$id <- rownames(edges)
          
          if (deggs_object[["padj_method"]] != "none") {
            edges$`p adj` <- formatC(edges$p.adj, format = "e", digits = 3)
            edges <- edges[order(edges$p.adj), ]
            edges <- edges[, which(colnames(edges) != "p.value")]
          } else {
            edges$`p value` <- formatC(edges$p.value, format = "e", digits = 3)
            edges <- edges[order(edges$p.value), ]
          }
          
          if (length(input$current_edges_selection) == 0) (
            DT::datatable(edges,
                          options = list(
                            lengthChange = FALSE, scrollX = TRUE,
                            columnDefs = list(list(
                              visible = FALSE,
                              targets = c(which(names(edges) == "id") - 1,
                                          which(names(edges) == sig_var) - 1)
                            ))
                          ),
                          rownames = FALSE
            )
          ) else (
            DT::datatable(
              edges[edges$id %in% input$current_edges_selection, ],
              options = list(
                lengthChange = FALSE, scrollX = TRUE,
                columnDefs = list(list(
                  visible = FALSE,
                  targets = c(which(names(edges) == "id") - 1,
                              which(names(edges) == sig_var) - 1)
                ))
              ),
              rownames = FALSE
            )
          )
        }
      }
    })
    
    # Network
    output$network <- visNetwork::renderVisNetwork({
      if (is.data.frame(category_networks[[input$category]])) {
        if (nrow(category_networks[[input$category]]) != 0) {
          edges <- category_networks[[input$category]]
          edges$id <- rownames(edges)

          # Set up tooltip
          prefix <- ifelse(deggs_object[["padj_method"]] == "none",
                           "P=", "Padj=")
          edges$title <- paste0(prefix, formatC(edges[, sig_var], format = "e",
                                                digits = 2))

          # Set up edges width
          # normalise p value between 0 and 1
          edges$width <- (edges[, sig_var] - min(edges[, sig_var])) /
            (max(edges[, sig_var]) - min(edges[, sig_var]))
          # invert values (and multiply by 4 to increase width)
          edges$width <- (1 - edges$width) * 4

          # Set up edges color
          if (multiOmic) {
            col <- my_palette(n = nlevels(edges$layer))
            edges$color <- col[edges$layer]  # edges$layer is factor
            legend <- levels(edges$layer)
          } else {
            col <- c("#00008B", "#9A9A9A")
            edges$color <- ifelse(edges[, sig_var] < 0.05, "#00008B", "#9A9A9A")
            legend <- c("significant", " not significant")
          }

          # Slider
          edges <- edges[edges[, sig_var] < input$slider, ]
          if (nrow(edges) == 0) {
            nodes <- data.frame()
            edges <- data.frame()
            network.title <- paste0("No differential interaction with ",
                                    sig_var, "<", input$slider,
                                    ".<br> Try to increase the ",
                                    sig_var, " threshold.")
            visNetwork::visNetwork(nodes, edges, main = network.title)
          } else {
            nodes <- data.frame(
              "id" = unique(c(edges$from, edges$to)),
              "label" = unique(c(edges$from, edges$to)),
              "title" = unique(c(edges$from, edges$to))
            )

            visNetwork::visNetwork(nodes, edges) %>%
              visNetwork::visIgraphLayout(physics = TRUE,
                                          smooth = TRUE,
                                          type = "full") %>%
              visNetwork::visNodes(color = list("border" = 'white'),
                                   font = list("size" = 16)) %>%
              visNetwork::visEdges(arrows = "to",
                                   color = list("inherit" = FALSE)) %>%
              visNetwork::visLayout(randomSeed = 12) %>%
              visNetwork::visLegend(
                useGroups = FALSE,
                position = 'right',
                addEdges = data.frame(
                  color = col,
                  label = legend,
                  font.align = "top",
                  font.size = 12
                ),
                stepY = stepY_legend,
                width = legend.arrow.width,
                zoom = TRUE
              ) %>%
              visNetwork::visOptions(highlightNearest = TRUE) %>%
              visNetwork::visInteraction(hover = TRUE, tooltipDelay = 20) %>%
              visNetwork::visEvents(
                select = "function(data) {
                Shiny.onInputChange('current_nodes_selection', data.nodes);
                Shiny.onInputChange('current_edges_selection', data.edges);
                }"
              )
          }
        } else {
          network.title <- "No differential interaction active for this category"
          nodes <- data.frame()
          edges <- data.frame()
          visNetwork::visNetwork(nodes, edges, main = network.title)
        }
      } else {
        network.title <- "No differential interaction active for this category"
        nodes <- data.frame()
        edges <- data.frame()
        visNetwork::visNetwork(nodes, edges, main = network.title)
      }
    })

    # Table
    output$tbl <- DT::renderDT({
      edges()
    })
    
    # Side Plot
    output$edge_or_node_plot <- shiny::renderPlot({
      edges <- category_networks[[input$category]]
      try(
        if (is.null(input$current_nodes_selection) &
            length(input$current_edges_selection) == 1) {

          plot_regressions(
            deggs_object = deggs_object,
            gene_A = edges[input$current_edges_selection, "from"],
            gene_B = edges[input$current_edges_selection, "to"],
            assayDataName = ifelse(multiOmic, 
                   as.character(edges[input$current_edges_selection, "layer"]),
                   1)
          )
        } else {
          shiny::req(input$current_nodes_selection != "")
          node_boxplot(
            deggs_object = deggs_object,
            gene = input$current_nodes_selection,
            assayDataName = 1
          )
        }
        ,silent = TRUE)
    })
    
    # Highligh the searched node in the network
    shiny::observe({
      if (is.data.frame(category_networks[[input$category]])) {
        if (input$searchButton > 0) {
          shiny::isolate({
            edges <- category_networks[[input$category]]
            nodes <- data.frame(
              "id" = unique(c(edges$from, edges$to)),
              "label" = unique(c(edges$from, edges$to)),
              "title" = unique(c(edges$from, edges$to))
            )
            
            nodes_selection$current_node <- nodes[grep(input$searchText, 
                                                       nodes$label,
                                                       ignore.case = TRUE),
                                                  "id"]
            visNetwork::visNetworkProxy("network") %>%
              visNetwork::visSelectNodes(id = nodes_selection$current_node)
          })
        }
      }
    })
    
    shiny::observe({
      shiny::updateSliderInput(session,
                               inputId = "slider",
                               max = max(round(outVar()[, sig_var]+ 0.001,
                                               digits = 3))
      )
    })
  }
  
  ################# UI ########################
  ui <- shiny::fluidPage(
    shiny::titlePanel("Differential Networks"),
    shiny::sidebarLayout(
      shiny::sidebarPanel(
        width = 4, # this will leave more space for the network
        
        # dropdown menu for category
        shiny::selectInput(
          inputId = "category", label = "Category",
          choices = levels(deggs_object[["metadata"]]),
          selected = levels(deggs_object[["metadata"]])[1]
        ),
        
        # Slider
        shiny::sliderInput("slider",
                           label = ifelse(
                             deggs_object[["padj_method"]] == "none",
                             "P values",
                             "Adjusted P values"),
                           min = 0.01,
                           max = 10, # fake, it will be updated by updateSliderInput
                           value = 0.05, step = 0.01
        ),
        
        # Table
        shiny::tags$div(DT::DTOutput('tbl'), style = "font-size: 75%"),
        
        # Searchbox
        shinydashboard::sidebarSearchForm(
          textId = "searchText",
          buttonId = "searchButton",
          label = "Search node..."
        ),
        
        # Plots
        shiny::tags$div(shiny::plotOutput('edge_or_node_plot'))
      ),
      shiny::mainPanel(
        visNetwork::visNetworkOutput("network", height = "700px")
      )
    ),
    shiny::tags$script(shiny::HTML("$(function() {
              $('#searchText').keypress(function(e) {
              if (e.which == 13) {
                $('#searchButton').click();
              }
              });
    });"))
  )
  shiny::shinyApp(ui = ui, server = server)
}


