# Load results:
load("./FAbatchPaper/Results/CrossbatchpredResults.Rda")


# Combine results to a data frame:
###################################

results <- data.frame(method=sapply(CrossbatchpredResults, function(x) x$method), 
batchlearnind=sapply(CrossbatchpredResults, function(x) x$batchlearnind), 
error=sapply(CrossbatchpredResults, function(x) x$error), 
mcc=sapply(CrossbatchpredResults, function(x) x$mcc))

results$method <- as.character(results$method)
results$method[results$method=="standardize"] <- "stand"
results$method[results$method=="meancenter"] <- "meanc"
results$method <- factor(results$method, levels=c("none", 
  "fabatch", "combat", "fsvafast", "fsvaexact", "meanc", 
  "stand", "ratiog", "ratioa")) 
results$batchlearnind[results$batchlearnind==1] <- "Training on the first batch"
results$batchlearnind[results$batchlearnind==2] <- "Training on the second batch"




# Plot of the MCC-values:
##########################

library("ggplot2")

p <- ggplot(data=results, aes(x=method, y=mcc)) + theme(axis.title.x=element_blank(), axis.text=element_text(size=11),
    axis.title=element_text(size=11), strip.text.x = element_text(size = 11), axis.text.x=element_text(angle=45, vjust=1, hjust=1)) +
  geom_point() + labs(y="MCC") + facet_wrap(~ batchlearnind, nrow=2, ncol=1) +
  scale_x_discrete(labels=gsub("Standardize", "Stand.", gsub("Meancenter", "Meanc.", levels(results$method))))
p

ggsave(file="./FAbatchPaper/Results/Figure3.pdf", width=7, height=6)





# PCA plots:
#############

# Auxiliary function - makse a single PCA plot (used below):

plotpc <- function(Xadj, y, batch, refbatch = 1) {

xpr <- prcomp(Xadj, scale. = FALSE)
PCs <- predict(xpr)[,1:2]

resultsPC <- data.frame(batch=factor(c(batch[batch==1], batch[batch==2]), levels=1:2), 
  y=factor(c(y[batch==1], y[batch==2]), levels=1:2), pc1=PCs[,1], pc2=PCs[,2])


batches <- unique(resultsPC$batch)

xall <- yall <- zall <- batchall <- c()

for(j in seq(along=batches)) {

pc1temp <- resultsPC$pc1[resultsPC$batch==batches[j]]
pc2temp <- resultsPC$pc2[resultsPC$batch==batches[j]]
dens2 <- kde2d(pc1temp, pc2temp, lims=c(min(pc1temp)-sd(pc1temp), max(pc1temp)+sd(pc1temp), min(pc2temp)-
  sd(pc2temp), max(pc2temp)+sd(pc2temp)))

densdf <- expand.grid(dens2$x, dens2$y)[,2:1]
densdf$z <- as.vector(dens2$z)

names(densdf) <- c("y", "x", "z")

xall <- c(xall, densdf$x)
yall <- c(yall, densdf$y)
zall <- c(zall, densdf$z)
batchall <- c(batchall, rep(batches[j], nrow(densdf)))

}

densdfall <- data.frame(x=xall, y=yall, z=zall, batch=batchall)

if(refbatch==2) {

densdfall$batch <- ifelse(densdfall$batch==1, 2, 1)
resultsPC$batch <- ifelse(resultsPC$batch==1, 2, 1)

}

p <- ggplot(data=resultsPC, aes(x=pc1, y=pc2)) + 
  scale_color_manual(values = c("black", "red")) +
  stat_contour(data=densdfall[densdfall$batch==1,], aes(x=x, y=y, z=z), size=0.8, colour="black", bins=5) + 
  geom_text(data=resultsPC[resultsPC$batch==1,], aes(label=as.character(as.numeric(y))), colour="black", fontface="bold", size=4) +
  stat_contour(data=densdfall[densdfall$batch==2,], aes(x=x, y=y, z=z), size=0.5, colour="red", bins=5) + 
  geom_text(data=resultsPC[resultsPC$batch==2,], aes(label=as.character(as.numeric(y))), colour="red", fontface="plain", size=4) +
  labs(x="PC 1", y="PC 2") + 
  geom_point(data=data.frame(x=sapply(batches, function(x) mean(resultsPC$pc1[resultsPC$batch==x])), 
    y=sapply(batches, function(x) mean(resultsPC$pc2[resultsPC$batch==x])), batch=batches), 
    aes(x=x, y=y), shape=18, size=5, colour="black") + 
  geom_point(data=data.frame(x=sapply(batches, function(x) mean(resultsPC$pc1[resultsPC$batch==x])), 
    y=sapply(batches, function(x) mean(resultsPC$pc2[resultsPC$batch==x])), batch=batches), 
    aes(x=x, y=y, colour=factor(batch)), shape=18, size=3.5) + 
  theme(legend.position="none",
    axis.title=element_text(size=12), axis.ticks=element_blank(),
    axis.text=element_blank(), axis.title.x=element_text(vjust=1.8), axis.title.y=element_text(vjust=-0.4))

return(p)

}



# Auxiliary function - performs batch effect adjustment on one batch using a 
# specified method and then addon batch effect adjustment on the other:

adjustaddon <- function(X, y, batch, method, refbatch = 1) {

require("bapred")

set.seed(1234)

testbatch <- setdiff(c(1,2), refbatch)

if(!(method %in% c("fsvafast", "fsvaexact"))) {
params <- ba(x=X[batch==refbatch,], y=y[batch==refbatch], batch=factor(rep(1, sum(batch==refbatch))), method = method)
}
else {
if(method=="fsvafast")
  params <- ba(x=X[batch==refbatch,], y=y[batch==refbatch], batch=factor(rep(1, sum(batch==refbatch))), method = "sva", algorithm = "fast")
if(method=="fsvaexact")
  params <- ba(x=X[batch==refbatch,], y=y[batch==refbatch], batch=factor(rep(1, sum(batch==refbatch))), method = "sva", algorithm = "exact")
}

Xtest <- baaddon(params=params, x=X[batch==testbatch,], batch=factor(rep(1, sum(batch==testbatch))))

if(refbatch==1)
  Xadj <- rbind(params$xadj, Xtest)
if(refbatch==2)
  Xadj <- rbind(Xtest, params$xadj)

return(Xadj)

}



# Load the dataset 'IUGRTranscr'.
# NOTE: The Rda-file loaded here is NOT included in the electronic appendix. To obtain it,
# one has to run the R-script IUGRTranscr_preparationinfos.R' found in the folder
# 'FAbatchPaper/Datasets/PreparationScripts'.

load("./FAbatchPaper/Datasets/ProcessedData/IUGRTranscr.Rda")


# A few variables are constant in the first batch --> exclude these:
X <- X[,apply(X[batch==1,], 2, sd)!=0]



# batch effect adjustment methods to use:
methods <- c("none", "fabatch", "combat", "fsvafast", "fsvaexact", "meancenter", 
  "standardize", "ratiog", "ratioa")
# short names of the above:
methodsshortnames <- c("none", "fabatch", "combat", "fsvafast", "fsvaexact", "meanc", 
  "stand", "ratiog", "ratioa")

  
  
# Make all subplots and store them in a list:

library("grid")

plist <- list()

for(i in seq(along=methods)) {
  cat(methods[i], "\n")
  plist[[i]] <- list()
  for(j in 1:2) {
    Xadj <- adjustaddon(X, y, batch, method = methods[i], refbatch = j)
    plist[[i]][[j]] <- plotpc(Xadj, y, batch, refbatch = j) + theme(plot.margin = unit(c(0.1,0.3,0.5,0.1), "cm"))
  }
  names(plist[[i]]) <- c("Training on first batch", "Training on second batch")
}

names(plist) <- methods



# For each method combine the two corresponding subplots to one plot:

library("gridExtra")

groblist <- list()
library("grid")
for(i in seq(along=methods))
  groblist[[i]] <- arrangeGrob(plist[[i]][[1]], plist[[i]][[2]], ncol=2, nrow=1, top=textGrob(methodsshortnames[i], gp=gpar(fontsize=14)))

  
  
# Arrange all plots into one single figure:  
  
grobframe <- marrangeGrob(groblist, ncol=2, nrow=5, top=element_blank())

grobframe

ggsave <- ggplot2::ggsave; body(ggsave) <- body(ggplot2::ggsave)[-2]


# Save the figure:

ggsave(file="./FAbatchPaper/Results/Figure4.pdf", plot=grobframe, width=9.1, height=11)
