################################################################################
# Function which:
# - calculates the correlations between drug response profiles
# - produces correlation matrix (either symmetric or 2 combined*)
#
# * for example, for two sets of samples coming from patients with different
#   diagnosis
################################################################################

makeCorrHeatmap = function(mt, mt2=NA, colsc, concNo="one",
                           ctab=BloodCancerMultiOmics2017::conctab,
                           dtab=BloodCancerMultiOmics2017::drugs) {
  
  # quiets concerns of R CMD check "no visible binding for global variable"
  NameX=NULL; NameY=NULL; Measure=NULL; x=NULL; y=NULL; xend=NULL;
  yend=NULL
  
  # Function which reorders the clustered matrix based on
  # computed singular value decomposition.
  callback = function(hc, mat){
    sv = svd(t(mat))$v[,1]
    dend = reorder(as.dendrogram(hc), wts = sv)
    rv <- as.hclust(dend)
    rv
  }
  
  
  # cluster the drugs
  mt = mt[apply(mt, 1, sd)!=0, ]
  mt = cor(t(mt))
  hc = hclust(as.dist(1-mt))
  hc = callback(hc, mt)
  ord = rev(hc$label[hc$order])
  
  # prepare data in LF
  mt = meltWholeDF(mt)
  mt$NameX = factor(mt$Y, levels=rev(ord))
  if(concNo=="one") {
    mt$NameY = factor(dtab[mt$X,"name"], levels=dtab[ord,"name"])
  } else if(concNo=="all") {
    mt$NameY = factor(giveDrugLabel(mt$X, ctab, dtab),
                      levels=giveDrugLabel(ord, ctab, dtab))
  }
  
  # take care of the secoond heatmap if present
  if(!is.na(mt2)[1]) {
    # obtain the correlation coefficients
    mt2 = mt2[apply(mt2, 1, sd)!=0, ]
    mt2 = cor(t(mt2))
    
    # order the matrix in the same way as mt is ordered
    mt2 = mt2[levels(mt$NameX), levels(mt$NameX)]
    # put 0 in lower triangle of matrix, include the diagonal
    mt2[upper.tri(mt2, diag=TRUE)] = NA
    
    mt2 = meltWholeDF(mt2)
    # remove NA rows
    mt2 = mt2[!is.na(mt2$Measure),]
    # insert top half of matrix data to plotting data frame
    idx = match(paste(mt2$X, mt2$Y, sep="."), paste(mt$X, mt$Y, sep="."))
    mt[idx, "Measure"] = mt2$Measure
  }
  
  # color the drug names on the axis
  pathColor = pathColor
  pathColor["Other"] = "black"
  if(concNo=="one") {
    drcol = unname(pathColor[match(dtab[rev(levels(mt$NameX)),"pathway"],
                                   names(pathColor))])
  } else if(concNo=="all") {
    drcol = unname(pathColor[match(dtab[rev(substring(levels(mt$NameX), 1, 5)),
                                        "pathway"], names(pathColor))])
  }
  
  # construct the main plot
  mainPlot = ggplot() +
    geom_tile(data=mt, aes(x=NameX, y=NameY, fill=Measure)) +
    theme_bw() +
    theme(axis.text.y = element_text(size=14, colour=drcol),
          panel.background=element_rect(fill=NA),
          legend.title=element_text(size=14),
          legend.text=element_text(size=14)) +
    scale_fill_gradientn(name="Correlation coefficient", limits=c(-1,1),
                         colours=colsc) +
    xlab("") + ylab("") + coord_fixed()
  
  # construct the dendrogram
  dhc = as.dendrogram(hc)
  ddata = dendro_data(dhc, type = "rectangle")
  ddata$segments = ddata$segments - 0.5
  nXY = nrow(ddata$labels)
  dendPlot = ggplot(segment(ddata)) +
    geom_segment(aes(x=x, y=y, xend=xend, yend=yend)) +
    theme_dendro() +
    scale_x_continuous(limits=c(0, nXY), expand = c(0,0)) +
    scale_y_reverse()
  
  # construct annotation for drug labels
  annoPlot = ggplot(data.frame(x=1, y=factor(names(pathColor),
                                             levels=names(pathColor))),
                    aes(x, y, fill=y)) +
    geom_tile() +
    scale_fill_manual("Targeted pathway", values=pathColor) +
    theme(legend.title=element_text(size=14), legend.text=element_text(size=14))
  
  # construct gtable
  wdths = c(3, 0.22*nXY, 0.1)
  hghts = c(0.1, 0.22*nXY, 1)
  
  gt = gtable(widths=unit(wdths, "in"), heights=unit(hghts, "in"))
  
  # make grobs
  heat = ggplotGrob(mainPlot)
  denX = ggplotGrob(dendPlot)
  anno = ggplotGrob(annoPlot)
  
  # fill in the gtable
  gt = gtable_add_grob(gt, gtable_filter(heat, "panel"), 2, 2)
  gt = gtable_add_grob(gt, denX$grobs[[whichInGrob(denX, "panel")]], 3, 2)
  gt = gtable_add_grob(gt, heat$grobs[[whichInGrob(heat, "axis-l")]], 2, 1)
  
  # gtable with legends
  wdthsl = c(3, 3)
  hghtsl = c(2.5)
  gtl = gtable(widths=unit(wdthsl, "in"), heights=unit(hghtsl, "in"))
  gtl = gtable_add_grob(gtl, gtable_filter(heat, "guide-box"), 1, 1)
  gtl = gtable_add_grob(gtl, gtable_filter(anno, "guide-box"), 1, 2)
  
  return(list("figure"=list(width=sum(wdths), height=sum(hghts), plot=gt),
              "legend"=list(width=sum(wdthsl), height=sum(hghtsl), plot=gtl)))
}