# Script contains helper functions for drug response prediction

# label the drugs with proper u
giveDrugLabel3 = function(drid, dtab=BloodCancerMultiOmics2017::drugs,
                          ctab=BloodCancerMultiOmics2017::conctab) {
  vapply(strsplit(drid, "_"), function(x) {
    k <- paste(x[1:2], collapse="_")
    paste0(dtab[k, "name"], " ",
           switch(x[3], "1:5"="c1:5", "4:5"="c4:5",
                  paste0(ctab[k, as.integer(x[3])], " \u00B5","M")))
  }, character(1))
}

# select features from the lpd object to use as response and predictors in the model
featureSelectionForLasso = function(objective, predictors, lpd) {
  
  # quiets concerns of R CMD check "no visible binding for global variable"
  dataType=NULL
  
  dimobj = dimnames(objective)
  objectiveName = dimobj[[1]][1]
  
  #design matix
  fx = t(Biobase::exprs(lpd)[predictors, ])

    # check the objective
  stopifnot(identical(rownames(fx), dimobj[[2]]))
  fy = objective[,,drop=TRUE]
  
  # select predictors
  fx = fx[, (fData(lpd)[predictors, "type"] %in% c("Methylation_Cluster",
                                                   "IGHV","gen", "pretreat")),
          drop=FALSE]
  
  #  make sure that there are no NAs in the table
  stopifnot(all(!is.na(fx)))
  
  # print table with numer of predictors from each group
  # message("Objective: ", objectiveName, "\n")
  prefix = paste(ifelse(dataType %in% unique(fData(lpd)[colnames(fx), "type"]),
                        names(dataType), "_"), collapse="")
  
  n = length(unique(fy))
  family = "gaussian"
  return(list(fx=fx, fy=fy, objectiveName=objectiveName, prefix=prefix,
              family=family))
}

# plotting the predictions
plotPredictions = function(fx, fy, objective, pred, coeffs, lpd, nm, lim,
                           objectiveName, colors) {

  # quiets concerns of R CMD check "no visible binding for global variable"
  X=NULL; Y=NULL; Measure=NULL
                           
  # PREPARE DATA FOR PLOTTING
  
  stopifnot( identical(dim(pred), c(length(fy), 1L)), identical(rownames(pred),
                                                                names(fy)) )
  ordy <- order(fy, pred[, 1])
  # design matrix of selected predictors (unscaled)
  mat  <- t(fx[ordy, names(coeffs), drop=FALSE])
  # human-readable names where available & drugs with concentrations
  nicename = names(coeffs) %>% `names<-`(names(coeffs))
  idx = grepl("D_", nicename)
  nicename[idx] = giveDrugLabel3(nicename[idx])
  nicename[!idx] = fData(lpd)[names(nicename)[!idx], "name"] %>%
    `names<-`(names(nicename)[!idx])
  nicename = ifelse(!is.na(nicename) & (nicename!=""), nicename, names(nicename))
  nicename = gsub("Methylation_Cluster", "Methylation cluster", nicename)
  rownames(mat) = nicename[rownames(mat)]
  # prepare plot title
  title=nm
  
  # CREATE INDIVIDUAL PARTS OF THE FIGURE  
  
  # bar plot
  stopifnot(all(coeffs<lim & coeffs>-lim))
  part1df = data.frame(coeffs,
                       nm=factor(names(coeffs), levels=names(rev(coeffs))))
  part1df$col= ifelse(rownames(part1df)=="IGHV", "I",
                      ifelse(rownames(part1df)=="Methylation_Cluster", "M",
                             ifelse(rownames(part1df)=="Pretreatment", "P","G")))
  
  part1 = ggplot(data=part1df, aes(x=nm, y=coeffs, fill=col)) +
    geom_bar(stat="identity", colour="black", position = "identity",
             width=.66, size=0.2) +theme_bw() +
    geom_hline(yintercept=0, size=0.3) + scale_x_discrete(expand=c(0,0.5)) +
    scale_y_continuous(expand=c(0,0)) + coord_flip(ylim=c(-lim,lim)) +
    theme(panel.grid.major=element_blank(),
          panel.background=element_blank(),
          panel.grid.minor = element_blank(),
          axis.text=element_text(size=8),
          panel.border=element_blank()) +
    xlab("") + ylab("Model Coefficients") +
    geom_vline(xintercept=c(0.5), color="black", size=0.6)+
    scale_fill_manual(c("M", "I", "G", "P"),
                      values=c(M=colors[["M"]][2],
                               I=colors[["I"]],
                               G=colors[["G"]],
                               P=colors[["P"]]))
  
  
  # heat map
  # mat contains selected predictors with status for each patient
  # (e.g. 0/1 for mutations and IGHV, 0/0.5/1 for M)
  # to assign colors Gosia added 5 to meth and 2 to IGHV values,resuting in
  # 5 LP, 5.5 IP, 6 HP, 2 unmut IGHV or mut, 4 IGHV mut, 3 mut, 7 pre-treatment
  idx = grep("Methylation cluster", rownames(mat))
  mat[idx,] = mat[idx,]+5
  ## ighv
  idx = grep("IGHV", rownames(mat))
  mat[idx,] = (mat[idx,]*2)+2
  ## gene
  rnm = sapply(rownames(mat), function(nm) strsplit(nm," ")[[1]][1])
  idx = rownames(fData(lpd))[fData(lpd)$type=="gen" & rownames(lpd) %in% rnm]
  idx = match(idx, rnm)
  mat[idx,] = mat[idx,]+2
  ## pretreat
  idx = grep("Pretreatment", rownames(mat))
  mat[idx,] = ifelse(mat[idx,]==0,2,7)
  
  mat2 = meltWholeDF(mat)
  mat2$Measure = factor(mat2$Measure, levels=sort(unique(mat2$Measure)))
  mat2$X = factor(mat2$X, levels=colnames(mat))
  mat2$Y = factor(mat2$Y, levels=rev(rownames(mat)))
  
  part2 = ggplot(mat2, aes(x=X, y=Y, fill=Measure)) +
    geom_tile() + theme_bw() +
    scale_fill_manual(name="Mutated",
                      values=c(`2`="gray96", `3`=paste0(colors["G"], "E5"),
                               `5`=colors[["M"]][1], `5.5`=colors[["M"]][2],
                               `6`=colors[["M"]][3], `7`=colors[["P"]],
                               `4`=paste0(colors["I"],"E5")), guide=FALSE) +
    scale_y_discrete(expand=c(0,0)) +
    theme(axis.text.y=element_text(hjust=0, size=14),
          axis.text.x=element_blank(),
          axis.ticks=element_blank(),
          panel.border=element_rect(colour="gainsboro"),
          plot.title=element_text(size=12),
          legend.title=element_text(size=12),
          legend.text=element_text(size=12),
          panel.background=element_blank(),
          panel.grid.major=element_blank(),
          panel.grid.minor=element_blank()) +
    xlab("patients") + ylab("") + ggtitle(title)
  if(length(levels(mat2$Y)) > 1) {
    part2 = part2 + geom_hline(yintercept=seq(1.5, length(levels(mat2$Y)), 1),
                               colour="gainsboro", size=0.2)
  }
  
  # scatter plot
  mat3 = fy[colnames(mat)]
  mat3 = data.frame(X=factor(names(mat3), levels=names(mat3)), Y=mat3*100)
  Yrange = range(mat3$Y)
  Yhangs = diff(Yrange)*0.05
  Ylims = c(Yrange[1]-Yhangs, Yrange[2]+Yhangs)
  Yran = diff(Yrange)
  Ybreaks = if(Yran<=15) 5 else if(Yran>15 & Yran<=30) 10 else if(Yran>30 & Yran<=40) 15 else if(Yran>40 & Yran<=60) 20 else 40
  
  part4 = ggplot(mat3, aes(x=X, y=Y)) +
    geom_point(shape=21, fill="dimgrey", colour="black", size=1.2) +
    theme_bw() +
    theme(panel.grid.minor=element_blank(),
          panel.grid.major.x=element_blank(),
          axis.title.x=element_blank(),
          axis.text.x=element_blank(),
          axis.ticks.x=element_blank(),
          axis.text.y=element_text(size=8),
          panel.border=element_rect(colour="dimgrey", size=0.1),
          panel.background=element_rect(fill="gray96")) +
    ylab("Viability [%]") +
    scale_y_continuous(limits=c(Yrange[1]-Yhangs, Yrange[2]+Yhangs),
                       breaks=seq(-200,200,Ybreaks))
  
  
  # MERGE PARTS OF THE FIGURE
  # construct the gtable
  wdths = c(4, 0.5, 0.05*ncol(mat), 0.05)
  hghts = c(0.3, 0.25*nrow(mat), 0.1, 0.4, 0.35)
  gt = gtable(widths=unit(wdths, "in"), heights=unit(hghts, "in"))
  ## make grobs
  gg1 = ggplotGrob(part1)
  gg2 = ggplotGrob(part2)
  gg4 = ggplotGrob(part4)
  ## fill in the gtable
  gt = gtable_add_grob(gt, gtable_filter(gg1, "panel"), 2, 1) # add boxplot
  gt = gtable_add_grob(gt, gtable_filter(gg2, "panel"), 2, 3) # add heatmap
  gt = gtable_add_grob(gt, gtable_filter(gg4, "panel"), 4, 3) # add scatterplot
  gt = gtable_add_grob(gt, gtable_filter(gg2, "xlab"), 5, 3)
  gt = gtable_add_grob(gt, gg2$grobs[[whichInGrob(gg2, "title")]], 1, 3)
  gt = gtable_add_grob(gt, gg4$grobs[[whichInGrob(gg4, "axis-l")]], 4, 2)
  gt = gtable_add_grob(gt, gg1$grobs[[whichInGrob(gg1, "axis-b")]], 3, 1)
  gt = gtable_add_grob(gt, gtable_filter(gg1, "xlab-b"), 4, 1)
  
  # now complicated: add the axis-l labels from gg2
  gia = which(gg2$layout$name == "axis-l")
  gga = gg2$grobs[[gia]]
  gax = gga$children[[2]]
  gax$widths = rev(gax$widths)
  gax$grobs = rev(gax$grobs)
  gt = gtable_add_cols(gt, gg2$widths[gg2$layout[gia, ]$l])
  gt = gtable_add_grob(gt, gax, 2, 5)
  
  wdth = convertUnit(gt$widths, "in", valueOnly=TRUE)[5]
  # add column on the right of appropriate width
  maxwdth = 2
  gt = gtable_add_cols(gt, unit(maxwdth-wdth, "in"))
  
  return(list(plot=gt, width=sum(wdths)+maxwdth, height=sum(hghts),
              name=make.names(objectiveName)))
}


# main function to fit Lasso and produce barplots to find genetic
# determinants of drug response
doLasso = function(objective, predictors, lpd,suffix="",
                  nm=NA, lim=0.21, ncv=10, nfolds=10, std=FALSE,
                  adaLasso = TRUE, colors) {
  
  #construct design and response matrix
  out = featureSelectionForLasso(objective, predictors, lpd)
  
  # for simplification
  fy = out[["fy"]]
  fx = out[["fx"]]
  family = out[["family"]]
  objectiveName = out[["objectiveName"]]
  prefix = out[["prefix"]]
  fxdim = dim(fx)
  print(sprintf("Prediction for: %s; #samples: %d; #features: %d",
                objectiveName, fxdim[1], fxdim[2]))
  
  # adaptive lasso for a more stable feature selection
  set.seed(19087)
  if(adaLasso){
    if(ncol(fx)>= nrow(fx)) {
      RidgeFit <- cv.glmnet(fx, fy, alpha = 0, standardize = std,
                            family = family, nfolds=10)
      # wRidge <- pmin(1/abs((coef(RidgeFit, s = RidgeFit$lambda.min))), 1e+300)
      wRidge <- 1/abs(coef(RidgeFit, s = RidgeFit$lambda.min))
      wRidge <- wRidge[-1]
      weights <- wRidge
    } else {
      lmFit <- lm(fy ~ fx)
      # wLM <- pmin(1/abs(coefficients(lmFit)[-1]), 1e+300)
      wLM <- 1/abs(coefficients(lmFit)[-1])
      weights <- wLM
    }
    excludedFeatures <- which(weights==Inf)
  } else {
    weights <- rep(1, ncol(fx))
    excludedFeatures <- NULL
  }
  
  #perform repeated cross-validation to find an optimal penalisatio
  #parameter minimizing the cross-validated MSE
  cv.out <- cvr.glmnet(Y=fy, X=fx, family=family,
                       alpha=1, nfolds=nfolds,
                       ncv=ncv, standardize=std,
                       exclude = excludedFeatures,
                       type.measure = "mse", penalty.factor = weights)
  
  # #fit Lasso model for optimal lambda                   
  fit = glmnet(y=fy, x=fx, family=family,
               alpha=1,
               standardize=std,
               exclude = excludedFeatures,
               lambda=cv.out$lambda, penalty.factor = weights)
  
  #get optimal lambda and corresponding predictors with coefficients
  lambda <- cv.out$lambda[which.min(cv.out$cvm)]
  coeffs <- coef(fit, lambda)
  coeffs_all <- coeffs
  coeffs <- as.vector(coeffs) %>%
    `names<-`(rownames(coeffs))   # cast from sparse matrix to ordinary vector
  coeffs <- coeffs[ coeffs!=0 ]
  
  # remove intercept term
  stopifnot(names(coeffs)[1]=="(Intercept)")
  if (length(coeffs) > 1) {
    coeffs <- coeffs[-1]
  } else {
    print("No (0) predictors for given parameters!")
    return(0)
  }
  coeffs <- sort(coeffs)
  
  # Residuals in the model
  pred <- predict(fit, newx = fx, s = lambda, type = "response")
  residuals <- pred[,1]-fy
  
  plot = plotPredictions(fx, fy, objective, pred, coeffs, lpd, nm, lim,
                         objectiveName, colors)
  
  return(list(residuals=residuals, coeffs=coeffs_all, plot=plot))
}

# Make list of predictors for the given lpd
makeListOfPredictors = function(lpd) {
  return(list(
    predictorsM = rownames(fData(lpd))[fData(lpd)$type=="Methylation_Cluster"],
    predictorsG = rownames(fData(lpd))[fData(lpd)$type=="gen"],
    predictorsI = rownames(fData(lpd))[fData(lpd)$type=="IGHV"],
    predictorsP = rownames(fData(lpd))[fData(lpd)$type=="pretreat"]
  ))
}


# Pre-process data: explore what is available & feature selection
prepareLPD = function(lpd, minNumSamplesPerGroup, withMC=TRUE) {
  
  # PRETREATMENT
  # update the expression set by adding row about pretreatment
  pretreated <- t(matrix(ifelse(
    BloodCancerMultiOmics2017::patmeta[colnames(lpd),
    "IC50beforeTreatment"], 0, 1),
    dimnames=list(colnames(lpd), "Pretreatment"))) 
  fdata_pretreat <- data.frame(name=NA, type="pretreat", id=NA, subtype=NA,
                               row.names="Pretreatment")
  lpd <- ExpressionSet(assayData=rbind(Biobase::exprs(lpd), pretreated),
                          phenoData=new("AnnotatedDataFrame", data=pData(lpd)),
                          featureData=new("AnnotatedDataFrame",
                                          rbind(fData(lpd), fdata_pretreat)))
  # METHYLATION
  Biobase::exprs(lpd)[fData(lpd)$type=="Methylation_Cluster",] =
    Biobase::exprs(lpd)[fData(lpd)$type=="Methylation_Cluster",]/2
  
  # IGHV: changing name from Uppsala to IGHV
  rownames(lpd)[which(fData(lpd)$type=="IGHV")] = "IGHV"
  
  # GENETICS
  # changing name od del13q any to del13q & remove the rest of del13q (mono, single)
  idx = which(rownames(lpd) %in% c("del13q14_bi", "del13q14_mono"))
  lpd = lpd[-idx,]
  rownames(lpd)[which(rownames(lpd)=="del13q14_any")] = "del13q14"
  
  # remove CHROMOTHRYPSIS
  if("Chromothripsis" %in% rownames(lpd))
    lpd = lpd[-which(rownames(lpd)=="Chromothripsis"),]
  
  # SELECT GOOD SAMPLES
  idx = !is.na(Biobase::exprs(lpd)["IGHV",])
  if(withMC) idx = idx & !is.na(Biobase::exprs(lpd)["Methylation_Cluster",])
  # cut out the data to have information about Methylation_Cluster and IGHV for all samples
  lpd = lpd[, idx]
  # for the genetics - remove the genes which do not have enough samples
  which2remove = names(
    which(!apply(Biobase::exprs(lpd)[rownames(lpd)[fData(lpd)$type %in%
                                            c("gen")],], 1, function(cl) {
    if(all(is.na(cl))) return(FALSE)
    if(sum(is.na(cl)) >= 0.1*length(cl)) return(FALSE)
    tmp = table(cl)
    return(length(tmp)==2 & all(tmp>=minNumSamplesPerGroup))
  })))
  lpd = lpd[-match(which2remove, rownames(lpd)),]
  
  # for the ones with NA put 0 instead
  featOther = rownames(lpd)[fData(lpd)$type %in%
                              c("IGHV", "Methylation_Cluster", "gen", "viab",
                                "pretreat")]
  tmp = Biobase::exprs(lpd)[featOther,]
  tmp[is.na(tmp)] = 0
  Biobase::exprs(lpd)[featOther,] = tmp
  
  return(lpd)
}


# wrapper functions to do Lasso model fitting, plotting and prediction
makePredictions = function(drs, frq, lpd, predictorList, lim, std=FALSE,
                           adaLasso = TRUE, colors) {
                           
  res = lapply(names(drs), function(typ) {
    setNames(lapply(drs[[typ]], function(dr) {
      if(typ=="1:5")
        nm <- paste0(BloodCancerMultiOmics2017::drugs[dr, "name"],
                     " (average of all concentrations)")
        else if(typ=="4:5")
          nm <- paste0(BloodCancerMultiOmics2017::drugs[dr, "name"],
                       " (average of ", paste(
                       BloodCancerMultiOmics2017::conctab[dr,4:5]*1000,
                         collapse = " and "), " nM)")
      # G & I & M & P
      doLasso(Biobase::exprs(lpd)[grepl(dr, rownames(lpd)) &
                           fData(lpd)$subtype==typ,, drop=FALSE], 
              predictors=with(predictorList,
                              c(predictorsI, predictorsG, predictorsM,
                                predictorsP)), 
              lpd=lpd,
              suffix=paste0("_","th0", "_c",gsub(":","-",typ)),
              nm=nm, lim=lim, colors=colors)
    }), nm=drs[[typ]])
  })
  return(res)
}


# Function to plot the legends
makeLegends = function(legendFor, colors) {
  
  # quiets concerns of R CMD check "no visible binding for global variable"
  x=NULL; y=NULL
  
  # select the colors needed
  colors = colors[names(colors) %in% legendFor]
  
  nleg = length(colors)
  wdths = rep(2, length.out=nleg); hghts = c(2)
  gtl = gtable(widths=unit(wdths, "in"), heights=unit(hghts, "in"))
  n=1

  # M
  if("M" %in% names(colors)) {
    Mgg = ggplot(data=data.frame(x=1, y=factor(c("LP","IP","HP"),
                                               levels=c("LP","IP","HP"))),
                 aes(x=x, y=y, fill=y)) + geom_tile() +
      scale_fill_manual(name="Methylation cluster",
                        values=setNames(colors[["M"]], nm=c("LP","IP","HP"))) +
      theme(legend.title=element_text(size=12),
            legend.text=element_text(size=12))
    gtl = gtable_add_grob(gtl, gtable_filter(ggplotGrob(Mgg), "guide-box"), 1, n)
    n = n+1
  }
  
  # I
  if("I" %in% names(colors)) {
    Igg = ggplot(data=data.frame(x=1,
                                 y=factor(c("unmutated","mutated"),
                                          levels=c("unmutated","mutated"))),
                 aes(x=x, y=y, fill=y)) + geom_tile() +
      scale_fill_manual(name="IGHV",
                        values=setNames(c("gray96",
                                          paste0(colors["I"], c("E5"))),
                                        nm=c("unmutated","mutated"))) +
      theme(legend.title=element_text(size=12),
            legend.text=element_text(size=12))
    gtl = gtable_add_grob(gtl, gtable_filter(ggplotGrob(Igg), "guide-box"), 1, n)
    n = n+1
  }
  
  # G
  if("G" %in% names(colors)) {
    Ggg = ggplot(data=data.frame(x=1,
                                 y=factor(c("wild type","mutated"),
                                          levels=c("wild type","mutated"))),
                 aes(x=x, y=y, fill=y)) + geom_tile() +
      scale_fill_manual(name="Gene",
                        values=setNames(c("gray96",
                                          paste0(colors["G"], c("E5"))),
                                        nm=c("wild type","mutated"))) +
      theme(legend.title=element_text(size=12),
            legend.text=element_text(size=12))
    gtl = gtable_add_grob(gtl, gtable_filter(ggplotGrob(Ggg), "guide-box"), 1, n)
    n = n+1
  }
  
  # P
  if("P" %in% names(colors)) {
    Pgg = ggplot(data=data.frame(x=1,
                                 y=factor(c("no","yes"),
                                          levels=c("no","yes"))),
                 aes(x=x, y=y, fill=y)) + geom_tile() +
      scale_fill_manual(name="Pretreatment",
                        values=setNames(c(colors[["P"]], "white"),
                                        nm=c("yes","no"))) +
      theme(legend.title=element_text(size=12),
            legend.text=element_text(size=12))
    gtl = gtable_add_grob(gtl, gtable_filter(ggplotGrob(Pgg), "guide-box"), 1, n)
    n = n+1
  }
  
  return(list(plot=gtl, width=sum(wdths), height=sum(hghts)))
}