library(data.table)
library(magrittr)
library(graph)
library(igraph)
library(RCytoscape)
library(GeneNet)
library(Hmisc)

fprintf <- function(...)cat(sprintf(...))

createSubGGM <- function(metabolites, covars, annotations, alpha=0.05, correction.method = "bonferroni", column.name="Sub.pathway", rm.unknown=TRUE){
  annotations<-orderAnnot(annotations, colnames(metabolites))
  
  metabolites<-scale(metabolites)
  subannotations<-copy(annotations)
  if(!"Fluid" %in% colnames(annotations)){
    subannotations[, name:=subannotations[, get(column.name)]]
  }else{
    subannotations[, name:=paste(subannotations[, get(column.name)], Fluid)]
  }
  
  mat<-as.matrix(eigenM(metabolites, subannotations[, name])$M)
  
  subannotations[, label:=subannotations[, get(column.name)]]
  subannotations<-subset(subannotations, !duplicated(name))
  
  if(rm.unknown){
    subannotations<-subannotations[Sub.pathway != "Unknown"]
    mat<-mat[,!grepl("Unknown", colnames(mat))]
  }
  
  return(createGGM(mat, covars = covars, annotations = subannotations, alpha = alpha))
}


orderAnnot<-function(annot_subnet_DT, names){
  
  newDT<-annot_subnet_DT[name==names[1]]
  for(i in 2:length(names)){
    newDT<-rbind(newDT, annot_subnet_DT[name==names[i]])
  }
  return(newDT)
}

createGGM <- function(metabolites, covars, annotations, alpha=0.05, correction.method = "bonferroni"){
  mat<-as.matrix(cbind(metabolites, covars))
  ggm<-ggm.estimate.pcor(mat, method = "dynamic")
  net_DT<-as.data.table(network.test.edges(ggm, direct=TRUE, fdr=TRUE, plot=FALSE))
  
  C_pear<-rcorr(mat, type = "pearson")
  colnames(C_pear$P)<-1:dim(C_pear$P)[1]
  rownames(C_pear$P)<-1:dim(C_pear$P)[1]
  Pear_DT<-data.table(node1= as.integer(rownames(C_pear$P)), C_pear$P)
  Pear_DT<-melt(Pear_DT, id.vars = "node1", variable.name = "node2", value.name = "ppear")
  Pear_DT[, node2 := as.integer(node2)]
  net_DT<-net_DT[Pear_DT, on=c("node1", "node2")]
  
  net_DT[, pval.adjust:=p.adjust(pval, method = correction.method)]
  net_DT[, ppear.adjust:=p.adjust(ppear, method = correction.method)]
  
  net_DT<-net_DT[pval.adjust <= alpha & ppear.adjust <= alpha]
  net_DT[, from:=gsub("\'", "", colnames(ggm))[node1]]
  net_DT[, to:=gsub("\'", "", colnames(ggm))[node2]]
  net_DT[, c("node1", "node2", "qval", "prob", "log.spvar", "pval.dir", "qval.dir", "prob.dir", "ppear", "ppear.adjust", "pval.adjust") := NULL]
  setcolorder(net_DT, c("from", "to", "pcor", "pval"))
  colnames(net_DT)<-c("from", "to", "Correlation", "p-Value")
  
  net_DT<-net_DT[to %in% annotations$name & from %in% annotations$name]
  
  net_graph<-graph_from_data_frame(d=net_DT, directed = F, vertices = annotations)
  
  return(net_graph)
}



# aggregated mean scores
# warning: does not scale, has to be done outside
# returns: matrix of aggregated values
aggMean = function(DF,G) {
  uG=unique(G)
  res=sapply(uG,function(g){apply(as.data.frame(DF[,G==g]),1,mean, na.rm=T  )})
  colnames(res)=uG
  res
}

# eigenvalues
# no data transformation inside function
# returns: list of two, M: matrix of eigenscores, expvar: full lists of explained variances
eigenM = function(DF,G) {
  uG=unique(G)
  res=sapply(uG,function(g){
    pca = prcomp(as.data.frame(DF[,G==g]))
    list(pc1=pca$x[,1],
         expvar=(pca$sdev)^2 / sum(pca$sdev^2) )
  }, simplify=F)
  # assemble matrix
  M = as.data.frame(sapply(res, function(x){x$pc1}))
  # assemble expvars
  expvar = sapply(res, function(x){x$expvar}, simplify=F)
  # return
  list(M=M,expvar=expvar)
}

identifyModules<-function(graph, metabolites, phenotype, covars, 
                          annotations,
                          merge.overlapping=FALSE, 
                          better.than.components= TRUE,
                          alpha=0.05,
                          for.sub.pathway=FALSE,
                          silent=FALSE){
  
  if(!is.data.table(metabolites)){
    if(is.data.frame(metabolites)){
      metabolites<-as.data.table(metabolites)
    }else if(is.matrix(metabolites)){
      metabolites<-data.table(sampleID= paste0("sample", 1:dim(metabolites)[1]),
                              metabolites)
      metabolites<-melt(data=metabolites, id.vars = "sampleID", variable.name = "name")
      if(!silent) print(metabolites)
    }
  }
  
  #Calculate the z-scores for metabolites
  metabolites[, z.score:=scale(value), by=.(name)]
  if(for.sub.pathway){
    
    if(!"Fluid" %in% colnames(annotations)){
      annotations[, Sub.pathway:=annotations[, "Sub.pathway"]]
    }else{
      annotations[, Sub.pathway:=paste(annotations[, get("Sub.pathway")], Fluid)]
    }
    
    metabolites<-metabolites[annotations[, .(name, Sub.pathway)], on="name"]
    metabolites[, name:= Sub.pathway]
    metabolites[, Sub.pathway := NULL]
  }
  
  if(dim(covars)[1]!=length(phenotype)){stop("Covars and Phenotype have a different amount of samples")}
  if(dim(covars)[1]!=length(unique(metabolites$sampleID))){
    stop("Metabolites and Covars have a different amount of samples")}
  if(!all(as.character(unique(metabolites$name)) %in% vertex_attr(graph, "name"))){
    warning("Not all of your metabolites are represented in the graph")
  }
  if(!all(vertex_attr(graph, "name") %in% as.character(unique(metabolites$name)))){
    warning("Not all of your nodes are represented in the metabolites data.table")
  }
  
  if(!silent)fprintf(paste("identifyModules could take some minutes", "\n"))
  
  modules_DT<-data.table(moduleID=integer(), module.score=numeric())
  node_DT<-data.table(moduleID=integer(), nodeID=integer(), name=character(), label=character(), 
                      order.added=integer(), score.after.adding=numeric())
  already_calculated<-data.table(key.value=character(), score=numeric())
  seed.scores<-c()
  for(v in V(graph)){
    if(!silent)fprintf(paste("Calculating best module for node ", v, " as seed", "\n"))
    module<-getOptimalModuleForSeed(graph, v, metabolites, phenotype, covars, alpha, 
                                    already_calculated, better.than.components)
    already_calculated<-module$already_calculated
    seed.scores<-c(seed.scores, module$seed.score)
    if(length(module$module) > 1 & module$module.score > -log(alpha/vcount(graph))){
      newID<-dim(modules_DT)[1]+1
      
      modules_DT<-rbind(modules_DT, data.table(moduleID=newID, module.score=module$module.score))
      node_DT<-rbind(node_DT, data.table(moduleID=newID, nodeID=module$module, 
                                         name=vertex_attr(graph, "name", module$module), 
                                         label=vertex_attr(graph, "label", module$module), 
                                         order.added=1:length(module$module), 
                                         score.after.adding=module$score.sequence))
      
    }
  }
  seed_scores_DT<-data.table(nodeID=V(graph), seed.score=seed.scores)
  #print(already_calculated)
  if(!silent)fprintf("Deleting duplicates\n")
  tmp_DT<-delete.duplicates(node_DT, modules_DT, graph)
  modules_DT<-unique(tmp_DT[, .(moduleID, module.score)])
  tmp_DT[,module.score := NULL]
  node_DT<-tmp_DT
  
  if(better.than.components){
    if(!silent)fprintf("Deleting less significant contained modules\n")
    tmp_DT<-delete.supergraphs(node_DT, modules_DT, graph)
    modules_DT<-unique(tmp_DT[, .(moduleID, module.score)])
    tmp_DT[,module.score := NULL]
    node_DT<-tmp_DT
  }
  
  if(merge.overlapping){
    if(!silent)fprintf("Merging overlapping modules\n")
    toMerge<-getToMerge(node_DT)
    node_DT<-toMerge[node_DT, on=.(moduleID)]
    node_DT[, moduleID := NULL]
    setnames(node_DT, c("moduleID", "nodeID", "name", "label", "order.added", "score.after.adding"))
    modules_DT<-getModulesFornodeDT(graph, metabolites, phenotype, covars, node_DT)
  }
  
  node_DT<-node_DT[seed_scores_DT, on=.(nodeID)]
  node_DT[, name:=vertex_attr(graph, "name", nodeID)]
  node_DT[, label:=vertex_attr(graph, "label", nodeID)]
  return(list(modules=modules_DT, vertices=node_DT))
}




delete.supergraphs<-function(node_DT, modules_DT, graph){
  DT<-modules_DT[node_DT, on=.(moduleID)]
  DT_max<-DT
  moduleIDs<-modules_DT$moduleID
  for(i in 1:(length(moduleIDs))){
    for(j in (i):length(moduleIDs)){
      
      if(i != j){
        v1<-sort(V(graph)[DT[ moduleID == moduleIDs[i],nodeID]])
        v2<-sort(V(graph)[DT[ moduleID == moduleIDs[j],nodeID]])
        intersect<-sort(intersection(v1, v2))
        
        if(length(intersect) == length(v1) && all(intersect == v1)){
          if(modules_DT[moduleID == moduleIDs[i], module.score] >= 
             modules_DT[moduleID == moduleIDs[j], module.score]){
            DT_max<-DT_max[moduleID !=moduleIDs[j]]
          }else{
            DT_max<-DT_max[moduleID !=moduleIDs[i]]
          }
        }else if(length(intersect) == length(v2) && all(intersect == v2)){
          if(modules_DT[moduleID ==moduleIDs[j], module.score] >= 
             modules_DT[moduleID ==moduleIDs[i], module.score]){
            DT_max<-DT_max[moduleID !=moduleIDs[i]]
          }else{
            DT_max<-DT_max[moduleID !=moduleIDs[j]]
          }
        }
      }
    }
  }
  return(DT_max)
}

delete.duplicates<-function(node_DT, modules_DT, graph){
  DT<-modules_DT[node_DT, on=.(moduleID)]
  DT_unique<-DT
  moduleIDs<-modules_DT$moduleID
  for(i in 1:(length(moduleIDs)-1)){
    for(j in (i+1):length(moduleIDs)){
      v1<-V(graph)[DT[ moduleID == moduleIDs[i],nodeID]]
      v2 <- V(graph)[DT[ moduleID == moduleIDs[j],nodeID]]
      if(length(difference(v1, v2)) == 0 & length(difference(v2, v1)) ==0){
        DT_unique<-DT_unique[moduleID !=moduleIDs[j]]
      }
    }
  }
  return(DT_unique)
}

getModulesFornodeDT<-function(graph, metabolites, phenotype, covars, node_DT){
  modules<-node_DT[, unique(moduleID)]
  module.scores<-c()
  for(module in modules){
    vertices<-node_DT[moduleID==module, nodeID]
    module.scores<-c(module.scores, calculateModScoreForPhenotype(graph, vertices,metabolites, phenotype, covars))
  }
  
  return(data.table(moduleID=modules, module.score=module.scores))
}

getToMerge<-function(node_DT){
  overlapping<-unique(node_DT[node_DT, on=.(nodeID), allow.cartesian=TRUE][,.(moduleID, i.moduleID)])
  overlap_graph<-graph_from_edgelist(as.matrix(overlapping), directed=F)
  cluster=clusters(overlap_graph)
  return(data.table(cluster=cluster$membership, moduleID=V(overlap_graph)))
}

getOptimalModuleForSeed<-function(graph, nodeNr, metabolites, phenotype, covars, 
                                  alpha=0.05, already_calculated, 
                                  better.than.components){
  module<-c(nodeNr)
  ##Calculate score for seed
  seed.score<-calculateModScoreForPhenotype(graph, nodeNr, metabolites, phenotype, covars)
  high.score<-seed.score
  old_module<-c()
  score.sequence<-c(seed.score)
  neighbor.index<-2
  highest.not.signifikant <- seed.score
  
  while(length(old_module) != length(module) & high.score == highest.not.signifikant){
    old_module<-module
    ##get all neighbors of current module
    neighbors<-getAllNeighbors(graph, module)
    for(neighbor in neighbors){
      new_module<-unique(c(old_module, neighbor))
      asString<-paste(new_module[order(new_module)], collapse = " ")
      
      neighbor.score<-calculateModScoreForPhenotype(graph, neighbor, metabolites, phenotype, covars)
      
      ##Lookup, if the module was already calculated earlier
      if(asString %in% already_calculated$key){
        current.score<-already_calculated[key.value == asString, score]
      }else{
        current.score<-calculateModScoreForPhenotype(graph, new_module, metabolites, phenotype, covars)
        already_calculated<-rbind(already_calculated, data.table(key.value = asString, score = current.score))
      }
      
      if(better.than.components){
        
        if(current.score > high.score & current.score > highest.not.signifikant){
          
          highest.not.signifikant<-current.score
          if(current.score > neighbor.score){
            score.sequence[neighbor.index]<-current.score
            high.score <- current.score
            module<-new_module
          }
          
        }
        
      }else{
        if(current.score > high.score){
          score.sequence[neighbor.index]<-current.score
          high.score <- current.score
          highest.not.signifikant<-high.score
          module<-new_module
        }
      }
      
    }
    neighbor.index <- neighbor.index+1
  }
  return(list(module=module, module.score=high.score, seed.score=seed.score, already_calculated=already_calculated, score.sequence=score.sequence))
}

getAllNeighbors<-function(graph, module){
  return(unique(unname(unlist(lapply(module, neighbors, graph=graph)))))
}

calculateModScoreForPhenotype<-function(graph, vertices, metabolites, phenotype, covars){
  metabolites<-metabolites[name %in% vertex_attr(graph, "name", vertices)]
  #Take the mean for all metabolites in each sample
  avgMetabolites<-metabolites[,.(mean=mean(z.score)),by=.(sampleID)]
  DT<-data.table(y= avgMetabolites[,mean], phenotype=phenotype, covars)
  lm.Module<-lm(formula= y ~ ., data=DT)
  sum<-summary(lm.Module)
  p_value<-sum$coefficients[2,"Pr(>|t|)"]
  return(-log(p_value))
}

writeModules<-function(mods, file.name){
  module_DT<-mods[[1]]
  node_DT<-mods[[2]]
  DT<-module_DT[node_DT, on=.(moduleID), nomatch=0]
  DT[, module.pval:= exp(-module.score)]
  DT[, node.pval:= exp(-seed.score)]
  DT[, module.pval.after.adding.node:= exp(-score.after.adding)]
  DT[,label := NULL]
  DT[,module.score := NULL]
  DT[,seed.score := NULL]
  DT[,score.after.adding:=NULL]
  DT<-DT[order(moduleID, order.added)]
  write.table(DT, file.name, sep = "\t", row.names = FALSE)
} 


drawModules<-function(graph, title="", node_DT, module_DT, draw.single.modules=TRUE, silent=FALSE){
  if(!silent)fprintf(paste("Cytoscape output could take some minutes...\n"))
  node_DT<-transformnodeDT(node_DT)
  node_DT[, is.significant:=node.pval < (0.05/vcount(graph))]
  node_DT[is.na(moduleID), moduleID:=0]
  modules<-module_DT$moduleID
  colors<-rainbow(length(modules))
  names(colors)<-modules
  drawNetworkWithModules(graph, title, node_DT, colors)
  
  if(draw.single.modules){
    l<-lapply(modules, drawModule, graph=graph, title=title, node_DT=node_DT[!is.na(moduleID)], module_DT=module_DT, colors=colors)
  }
}

drawNetworkWithModules<-function(graph, title, node_DT, colors, save.image=TRUE, save.gml=FALSE){
  vertices<-node_DT
  module_graph<-graph
  for(i in 1:dim(vertices)[1]){
    module_graph<-set_vertex_attr(module_graph, "label", vertices[i]$name, 
                                  paste0(
                                    toupper(substr(vertex_attr(graph, "Fluid" ,vertices[i]$name),0,1)), 
                                    "::", vertices[i]$label))
    module_graph<-set_vertex_attr(module_graph, "module.name", vertices[i]$name, vertices[i]$moduleID)
    module_graph<-set_vertex_attr(module_graph, "p.value", vertices[i]$name, vertices[i]$node.pval)
    module_graph<-set_vertex_attr(module_graph, "is.significant", vertices[i]$name, vertices[i]$is.significant)
  }
  
  g<-as_graphnel(module_graph)
  g<-initNodeAttribute (g, "ID", "char", "")
  g<-initNodeAttribute (g, "label", "char", "")
  g<-initNodeAttribute (g, "Super.pathway", "char", "")
  g<-initNodeAttribute (g, "Sub.pathway", "char", "")
  g<-initNodeAttribute (g, "Fluid", "char", "")
  g<-initNodeAttribute (g, "p.value", "numeric", "")
  g<-initNodeAttribute (g, "module.name", "char", "x")
  g<-initNodeAttribute (g, "is.significant", "char", "")
  g <- initEdgeAttribute (g, "Correlation", "numeric", 0.0)
  g <- initEdgeAttribute (g, "weight", "numeric", 0.0)
  g <- initEdgeAttribute (g, "p-Value", "numeric", 0.0)
  cy = CytoscapeConnection ()
  cw = new.CytoscapeWindow (title, g)
  displayGraph(cw)
  setNodeLabelRule (cw, "label")
  setNodeSizeRule(cw, "p.value", c(0.05/vcount(graph), 0.05/vcount(graph),0.05)*1.01, c(100,50,35,25,20), mode = "interpolate")
  setNodeShapeRule(cw, "is.significant", c("TRUE", "FALSE"), c("diamond", "ellipse"))
  setNodeColorRule(cw,"module.name" , names(colors), c(substr(colors, 0, 7)), mode="lookup", default.color='#D8D0CE')
  setWindowSize (cw, 600, 1200)
  layoutNetwork (cw, "kamada-kawai")
  
  if(save.image){
    saveImage(cw, paste0(getwd(), "/", title, ".png"), "png", 7.0)
  }
  
  if(save.gml){
    saveNetwork(cw, paste0(getwd(), "/", title, ".gml"), "gml")
  }
  return(cw)
}

drawModule<-function(moduleNR, graph, title="", node_DT, module_DT, colors, save.image=TRUE){
  
  vertices<-node_DT[moduleID==moduleNR]
  module_graph<-induced_subgraph(graph, vertices$nodeID)
  for(i in 1:dim(vertices)[1]){
    module_graph<-set_vertex_attr(module_graph, "label", vertices[i]$name, 
                                  paste0(
                                    toupper(substr(vertex_attr(graph, "Fluid" ,vertices[i]$name),0,1)), 
                                    "::", vertices[i]$label))
    module_graph<-set_vertex_attr(module_graph, "module.name", vertices[i]$name, moduleNR)
    module_graph<-set_vertex_attr(module_graph, "p.value", vertices[i]$name, vertices[i]$node.pval)
    module_graph<-set_vertex_attr(module_graph, "is.significant", vertices[i]$name, vertices[i]$is.significant)
  }
  g<-as_graphnel(module_graph)
  g<-initNodeAttribute (g, "ID", "char", "")
  g<-initNodeAttribute (g, "label", "char", "")
  g<-initNodeAttribute (g, "Super.pathway", "char", "")
  g<-initNodeAttribute (g, "Sub.pathway", "char", "")
  g<-initNodeAttribute (g, "Fluid", "char", "")
  g<-initNodeAttribute (g, "p.value", "numeric", "")
  g<-initNodeAttribute (g, "module.name", "char", "x")
  g<-initNodeAttribute (g, "is.significant", "char", "")
  g <- initEdgeAttribute (g, "Correlation", "numeric", 0.0)
  g <- initEdgeAttribute (g, "weight", "numeric", 0.0)
  g <- initEdgeAttribute (g, "p-Value", "numeric", 0.0)
  cy = CytoscapeConnection ()
  cw = new.CytoscapeWindow (paste0(title, "_", "module", moduleNR), g)
  displayGraph(cw)
  setNodeLabelRule (cw, "label")
  setNodeSizeRule(cw, "p.value", c(0.05/vcount(graph), 0.05/vcount(graph)*0.01,0.05), c(80,60,45,30,20), mode = "interpolate")
  setNodeShapeRule(cw, "is.significant", c("TRUE", "FALSE"), c("diamond", "ellipse"))
  setNodeColorRule(cw,"module.name" , names(colors), c(substr(colors, 0, 7)), mode="lookup", default.color='#D8D0CE')
  
  setWindowSize (cw, 1200, 1200)
  layoutNetwork (cw, "kamada-kawai")
  
  
  if(save.image){
    saveImage(cw, paste0(getwd(), "/", title, "_", "module", moduleNR, ".png"), "png", 1.0)
  }
  return(cw)
  
}

transformnodeDT<-function(node_DT){

  return(node_DT[, .(moduleID, nodeID, name, label, order.added, node.pval=exp(-seed.score), after.adding.pval=exp(-score.after.adding))])
}