diff --git a/R/ComBat.R b/R/ComBat.R
index 9d6e100..aed307c 100644
--- a/R/ComBat.R
+++ b/R/ComBat.R
@@ -7,7 +7,9 @@
 #'
 #' @param dat Genomic measure matrix (dimensions probe x sample) - for example, expression matrix
 #' @param batch {Batch covariate (only one batch allowed)}
-#' @param mod Model matrix for outcome of interest and other covariates besides batch
+#' @param dat_test (Optional) Independent test data of the same nature than dat -e.g. expression matrix- and with the same rows. Useful for machine learning applications.
+#' @param batch_test (Optional) Batch covariate of dat_test.
+#' @param mod (Optional) Model matrix for outcome of interest and other covariates besides batch
 #' @param par.prior (Optional) TRUE indicates parametric adjustments will be used, FALSE indicates non-parametric adjustments will be used
 #' @param prior.plots (Optional) TRUE give prior plots with black as a kernel estimate of the empirical batch effect density and red as the parametric
 #' @param mean.only (Optional) FALSE If TRUE ComBat only corrects the mean of the batch effect (no scale adjustment)
@@ -40,267 +42,339 @@
 #' # reference-batch version, with covariates
 #' combat_edata3 = ComBat(dat=edata, batch=batch, mod=mod, par.prior=TRUE, ref.batch=3)
 #'
+#' # Training and test
+#' te_idx <- c(1:4,9:11,18:22,47:50) #Test indices
+#' combat_edata4 = ComBat(dat=edata[,-te_idx], batch=batch[-te_idx],dat_test=edata[,te_idx], batch_test=batch[te_idx])
+#'
 #' @export
 #'
 
-ComBat <- function(dat, batch, mod = NULL, par.prior = TRUE, prior.plots = FALSE,
-                   mean.only = FALSE, ref.batch = NULL, BPPARAM = bpparam("SerialParam")) {
-    if(length(dim(batch))>1){
-      stop("This version of ComBat only allows one batch variable")
-    }  ## to be updated soon!  
-    
-    ## coerce dat into a matrix
-    dat <- as.matrix(dat)
-    
-    ## find genes with zero variance in any of the batches
-    batch <- as.factor(batch)
-    zero.rows.lst <- lapply(levels(batch), function(batch_level){
-      if(sum(batch==batch_level)>1){
-        return(which(apply(dat[, batch==batch_level], 1, function(x){var(x)==0})))
-      }else{
-        return(which(rep(1,3)==2))
-      }
-    })
-    zero.rows <- Reduce(union, zero.rows.lst)
-    keep.rows <- setdiff(1:nrow(dat), zero.rows)
-    
-    if (length(zero.rows) > 0) {
-      cat(sprintf("Found %d genes with uniform expression within a single batch (all zeros); these will not be adjusted for batch.\n", length(zero.rows)))
-      # keep a copy of the original data matrix and remove zero var rows
-      dat.orig <- dat
-      dat <- dat[keep.rows, ]
-    }
+
+
+
+ComBat <- function(dat, batch, dat_test=NULL, batch_test=NULL,
+                   mod = NULL, par.prior = TRUE, prior.plots = FALSE, mean.only = FALSE,
+                   ref.batch = NULL, BPPARAM = bpparam("SerialParam")) {
   
-    ## make batch a factor and make a set of indicators for batch
-    if(any(table(batch)==1)){mean.only=TRUE}
-    if(mean.only==TRUE){
-        message("Using the 'mean only' version of ComBat")
+  
+  if(length(dim(batch))>1||length(dim(batch_test))>1){
+    stop("This version of ComBat only allows one batch variable")
+  }  ## to be updated soon!  
+  
+  
+  ## coerce dat into a matrix
+  dat <- as.matrix(dat)
+  
+  ## find genes with zero variance in any of the batches
+  batch <- as.factor(batch)
+  zero.rows.lst <- lapply(levels(batch), function(batch_level){
+    if(sum(batch==batch_level)>1){
+      return(which(apply(dat[, batch==batch_level], 1, function(x){var(x)==0})))
+    }else{
+      return(which(rep(1,3)==2))
     }
-    
-    batchmod <- model.matrix(~-1+batch)  
-    if (!is.null(ref.batch)){
-        ## check for reference batch, check value, and make appropriate changes
-        if (!(ref.batch%in%levels(batch))) {
-            stop("reference level ref.batch is not one of the levels of the batch variable")
-        }
-        message("Using batch =",ref.batch, "as a reference batch (this batch won't change)")
-        ref <- which(levels(as.factor(batch))==ref.batch) # find the reference
-        batchmod[,ref] <- 1
-    } else {
-        ref <- NULL
+  })
+  zero.rows <- Reduce(union, zero.rows.lst)
+  keep.rows <- setdiff(1:nrow(dat), zero.rows)
+  
+  if (length(zero.rows) > 0) {
+    cat(sprintf("Found %d genes with uniform expression within a single batch (all zeros); these will not be adjusted for batch.\n", length(zero.rows)))
+    # keep a copy of the original data matrix and remove zero var rows
+    dat.orig <- dat
+    dat <- dat[keep.rows, ]
+  }
+  
+  ## make batch a factor and make a set of indicators for batch
+  if(any(table(batch)==1)){mean.only=TRUE}
+  if(mean.only==TRUE){
+    message("Using the 'mean only' version of ComBat")
+  }
+  
+  batchmod <- model.matrix(~-1+batch)  
+  if (!is.null(ref.batch)){
+    ## check for reference batch, check value, and make appropriate changes
+    if (!(ref.batch%in%levels(batch))) {
+      stop("reference level ref.batch is not one of the levels of the batch variable")
     }
-    message("Found", nlevels(batch), "batches")
+    message("Using batch =",ref.batch, "as a reference batch (this batch won't change)")
+    ref <- which(levels(as.factor(batch))==ref.batch) # find the reference
+    batchmod[,ref] <- 1
+  } else {
+    ref <- NULL
+  }
+  message("Found", nlevels(batch), "batches")
   
-    ## A few other characteristics on the batches
-    n.batch <- nlevels(batch)
-    batches <- list()
-    for (i in 1:n.batch) {
-        batches[[i]] <- which(batch == levels(batch)[i])
-    } # list of samples in each batch  
-    n.batches <- sapply(batches, length)
-    if(any(n.batches==1)){
-        mean.only=TRUE
-        message("Note: one batch has only one sample, setting mean.only=TRUE")
+  ## A few other characteristics on the batches
+  n.batch <- nlevels(batch)
+  batches <- list()
+  for (i in 1:n.batch)     batches[[i]] <- which(batch == levels(batch)[i])
+  # list of samples in each batch  
+  n.batches <- sapply(batches, length)
+  if(any(n.batches==1)){
+    mean.only=TRUE
+    message("Note: one batch has only one sample, setting mean.only=TRUE")
+  }
+  n.array <- sum(n.batches)
+  ## combine batch variable and covariates
+  design <- cbind(batchmod,mod)
+  
+  ## check for intercept in covariates, and drop if present
+  check <- apply(design, 2, function(x) all(x == 1))
+  if(!is.null(ref)){
+    check[ref] <- FALSE
+  } ## except don't throw away the reference batch indicator
+  design <- as.matrix(design[,!check])
+  
+  ## Number of covariates or covariate levels
+  message("Adjusting for", ncol(design)-ncol(batchmod), 'covariate(s) or covariate level(s)')
+  
+  ## Check if the design is confounded
+  if(qr(design)$rank < ncol(design)) {
+    ## if(ncol(design)<=(n.batch)){stop("Batch variables are redundant! Remove one or more of the batch variables so they are no longer confounded")}
+    if(ncol(design)==(n.batch+1)) {
+      stop("The covariate is confounded with batch! Remove the covariate and rerun ComBat")
     }
-    n.array <- sum(n.batches)
-    ## combine batch variable and covariates
-    design <- cbind(batchmod,mod)
-  
-    ## check for intercept in covariates, and drop if present
-    check <- apply(design, 2, function(x) all(x == 1))
-    if(!is.null(ref)){
-        check[ref] <- FALSE
-    } ## except don't throw away the reference batch indicator
-    design <- as.matrix(design[,!check])
-  
-    ## Number of covariates or covariate levels
-    message("Adjusting for", ncol(design)-ncol(batchmod), 'covariate(s) or covariate level(s)')
-  
-    ## Check if the design is confounded
-    if(qr(design)$rank < ncol(design)) {
-        ## if(ncol(design)<=(n.batch)){stop("Batch variables are redundant! Remove one or more of the batch variables so they are no longer confounded")}
-        if(ncol(design)==(n.batch+1)) {
-            stop("The covariate is confounded with batch! Remove the covariate and rerun ComBat")
-        }
-        if(ncol(design)>(n.batch+1)) {
-            if((qr(design[,-c(1:n.batch)])$rank<ncol(design[,-c(1:n.batch)]))){
-                stop('The covariates are confounded! Please remove one or more of the covariates so the design is not confounded')
-            } else {
-                stop("At least one covariate is confounded with batch! Please remove confounded covariates and rerun ComBat")
-            }
-        }
+    if(ncol(design)>(n.batch+1)) {
+      if((qr(design[,-c(1:n.batch)])$rank<ncol(design[,-c(1:n.batch)]))){
+        stop('The covariates are confounded! Please remove one or more of the covariates so the design is not confounded')
+      } else {
+        stop("At least one covariate is confounded with batch! Please remove confounded covariates and rerun ComBat")
+      }
     }
+  }
   
-    ## Check for missing values
-    NAs <- any(is.na(dat))
-    if(NAs){
-        message(c('Found',sum(is.na(dat)),'Missing Data Values'), sep=' ')}
-    ## print(dat[1:2,])
-  
-    ##Standardize Data across genes
-    message('Standardizing Data across genes')
-    if (!NAs){
-        B.hat <- solve(crossprod(design), tcrossprod(t(design), as.matrix(dat)))
-    } else { 
-        B.hat <- apply(dat, 1, Beta.NA, design) # FIXME
+  ## Check for missing values
+  NAs <- any(is.na(dat))
+  if(NAs){
+    message(c('Found',sum(is.na(dat)),'Missing Data Values'), sep=' ')}
+  ## print(dat[1:2,])
+  
+  ##Standardize Data across genes
+  message('Standardizing Data across genes')
+  if (!NAs){
+    B.hat <- solve(crossprod(design), tcrossprod(t(design), as.matrix(dat)))
+  } else { 
+    B.hat <- apply(dat, 1, Beta.NA, design) # FIXME
+  }
+  
+  ## change grand.mean for ref batch
+  if(!is.null(ref.batch)){
+    grand.mean <- t(B.hat[ref, ])
+  } else {
+    grand.mean <- crossprod(n.batches/n.array, B.hat[1:n.batch,])
+  }
+  
+  ## change var.pooled for ref batch
+  if (!NAs){
+    if(!is.null(ref.batch)) {
+      ref.dat <- dat[, batches[[ref]]]
+      var.pooled <- ((ref.dat-t(design[batches[[ref]], ] %*% B.hat))^2) %*% rep(1/n.batches[ref],n.batches[ref]) # FIXME
+    } else {
+      var.pooled <- ((dat-t(design %*% B.hat))^2) %*% rep(1/n.array,n.array) # FIXME
     }
-    
-    ## change grand.mean for ref batch
-    if(!is.null(ref.batch)){
-        grand.mean <- t(B.hat[ref, ])
+  } else {
+    if(!is.null(ref.batch)) {
+      ref.dat <- dat[, batches[[ref]]]
+      var.pooled <- rowVars(ref.dat-t(design[batches[[ref]], ]%*%B.hat), na.rm=TRUE)
     } else {
-        grand.mean <- crossprod(n.batches/n.array, B.hat[1:n.batch,])
+      var.pooled <- rowVars(dat-t(design %*% B.hat), na.rm=TRUE)
     }
+  }
+  
+  stand.mean <- t(grand.mean) %*% t(rep(1,n.array)) # FIXME
+  
+  if(!is.null(design)){
+    tmp <- design
+    tmp[,c(1:n.batch)] <- 0
+    stand.mean <- stand.mean+t(tmp %*% B.hat)
+  }  
+  
+  s.data <- (dat-stand.mean)/(sqrt(var.pooled) %*% t(rep(1,n.array))) # FIXME
+  
+  ##Get regression batch effect parameters
+  message("Fitting L/S model and finding priors")
+  batch.design <- design[, 1:n.batch]
   
-    ## change var.pooled for ref batch
-    if (!NAs){
-        if(!is.null(ref.batch)) {
-            ref.dat <- dat[, batches[[ref]]]
-            var.pooled <- ((ref.dat-t(design[batches[[ref]], ] %*% B.hat))^2) %*% rep(1/n.batches[ref],n.batches[ref]) # FIXME
-        } else {
-            var.pooled <- ((dat-t(design %*% B.hat))^2) %*% rep(1/n.array,n.array) # FIXME
-        }
+  if (!NAs){
+    gamma.hat <- solve(crossprod(batch.design), tcrossprod(t(batch.design),
+                                                           as.matrix(s.data)))
+  } else{
+    gamma.hat <- apply(s.data, 1, Beta.NA, batch.design) # FIXME
+  }
+  delta.hat <- NULL
+  for (i in batches){
+    if(mean.only==TRUE) {
+      delta.hat <- rbind(delta.hat,rep(1,nrow(s.data))) 
     } else {
-        if(!is.null(ref.batch)) {
-            ref.dat <- dat[, batches[[ref]]]
-            var.pooled <- rowVars(ref.dat-t(design[batches[[ref]], ]%*%B.hat), na.rm=TRUE)
-        } else {
-            var.pooled <- rowVars(dat-t(design %*% B.hat), na.rm=TRUE)
-        }
+      delta.hat <- rbind(delta.hat, rowVars(s.data[,i], na.rm=TRUE))
     }
+  }
   
-    stand.mean <- t(grand.mean) %*% t(rep(1,n.array)) # FIXME
-    if(!is.null(design)){
-        tmp <- design
-        tmp[,c(1:n.batch)] <- 0
-        stand.mean <- stand.mean+t(tmp %*% B.hat) #FIXME
-    }  
-    s.data <- (dat-stand.mean)/(sqrt(var.pooled) %*% t(rep(1,n.array))) # FIXME
+  ##Find Priors
+  gamma.bar <- rowMeans(gamma.hat)
+  t2 <- rowVars(gamma.hat)
   
-    ##Get regression batch effect parameters
-    message("Fitting L/S model and finding priors")
-    batch.design <- design[, 1:n.batch]
-    if (!NAs){
-      gamma.hat <- solve(crossprod(batch.design), tcrossprod(t(batch.design),
-                                                             as.matrix(s.data)))
-    } else{
-        gamma.hat <- apply(s.data, 1, Beta.NA, batch.design) # FIXME
-    }
-    delta.hat <- NULL
-    for (i in batches){
-        if(mean.only==TRUE) {
-            delta.hat <- rbind(delta.hat,rep(1,nrow(s.data))) 
-        } else {
-            delta.hat <- rbind(delta.hat, rowVars(s.data[,i], na.rm=TRUE))
-        }
-    }
+  a.prior <- apply(delta.hat, 1, aprior) # FIXME 
+  b.prior <- apply(delta.hat, 1, bprior) # FIXME
   
-    ##Find Priors
-    gamma.bar <- rowMeans(gamma.hat)
-    t2 <- rowVars(gamma.hat)
-    a.prior <- apply(delta.hat, 1, aprior) # FIXME 
-    b.prior <- apply(delta.hat, 1, bprior) # FIXME
+  ## Plot empirical and parametric priors
   
-    ## Plot empirical and parametric priors
-
-    if (prior.plots && par.prior) {
-        old_pars <- par(no.readonly = TRUE)
-        on.exit(par(old_pars))
-        par(mfrow=c(2,2))
-        
-        ## Top left
-        tmp <- density(gamma.hat[1,])
-        plot(tmp,  type='l', main=expression(paste("Density Plot of First Batch ",  hat(gamma))))
-        xx <- seq(min(tmp$x), max(tmp$x), length=100)
-        lines(xx,dnorm(xx,gamma.bar[1],sqrt(t2[1])), col=2)
-        
-        ## Top Right
-        qqnorm(gamma.hat[1,], main=expression(paste("Normal Q-Q Plot of First Batch ", hat(gamma))))
-        qqline(gamma.hat[1,], col=2)
-        
-        ## Bottom Left
-        tmp <- density(delta.hat[1,])
-        xx <- seq(min(tmp$x), max(tmp$x), length=100)
-        tmp1 <- list(x=xx, y=dinvgamma(xx, a.prior[1], b.prior[1]))
-        plot(tmp, typ="l", ylim=c(0, max(tmp$y, tmp1$y)),
-             main=expression(paste("Density Plot of First Batch ", hat(delta))))
-        lines(tmp1, col=2)
-
-        ## Bottom Right
-        invgam <- 1/qgamma(1-ppoints(ncol(delta.hat)), a.prior[1], b.prior[1])
-        qqplot(invgam, delta.hat[1,],
-               main=expression(paste("Inverse Gamma Q-Q Plot of First Batch ", hat(delta))),
-               ylab="Sample Quantiles", xlab="Theoretical Quantiles")
-        lines(c(0, max(invgam)), c(0, max(invgam)), col=2)
+  if (prior.plots && par.prior) {
+    old_pars <- par(no.readonly = TRUE)
+    on.exit(par(old_pars))
+    par(mfrow=c(2,2))
+    
+    ## Top left
+    tmp <- density(gamma.hat[1,])
+    plot(tmp,  type='l', main=expression(paste("Density Plot of First Batch ",  hat(gamma))))
+    xx <- seq(min(tmp$x), max(tmp$x), length=100)
+    lines(xx,dnorm(xx,gamma.bar[1],sqrt(t2[1])), col=2)
+    
+    ## Top Right
+    qqnorm(gamma.hat[1,], main=expression(paste("Normal Q-Q Plot of First Batch ", hat(gamma))))
+    qqline(gamma.hat[1,], col=2)
+    
+    ## Bottom Left
+    tmp <- density(delta.hat[1,])
+    xx <- seq(min(tmp$x), max(tmp$x), length=100)
+    tmp1 <- list(x=xx, y=dinvgamma(xx, a.prior[1], b.prior[1]))
+    plot(tmp, typ="l", ylim=c(0, max(tmp$y, tmp1$y)),
+         main=expression(paste("Density Plot of First Batch ", hat(delta))))
+    lines(tmp1, col=2)
+    
+    ## Bottom Right
+    invgam <- 1/qgamma(1-ppoints(ncol(delta.hat)), a.prior[1], b.prior[1])
+    qqplot(invgam, delta.hat[1,],
+           main=expression(paste("Inverse Gamma Q-Q Plot of First Batch ", hat(delta))),
+           ylab="Sample Quantiles", xlab="Theoretical Quantiles")
+    lines(c(0, max(invgam)), c(0, max(invgam)), col=2)
+  }
+  
+  ## Find EB batch adjustments
+  
+  gamma.star <- delta.star <- matrix(NA, nrow=n.batch, ncol=nrow(s.data))
+  if (par.prior) {
+    message("Finding parametric adjustments")
+    results <- bplapply(1:n.batch, function(i) {
+      if (mean.only) {
+        gamma.star <- postmean(gamma.hat[i,], gamma.bar[i], 1, 1, t2[i])
+        delta.star <- rep(1, nrow(s.data))
+      }
+      else {
+        temp <- it.sol(s.data[, batches[[i]]], gamma.hat[i, ],
+                       delta.hat[i, ], gamma.bar[i], t2[i], a.prior[i],
+                       b.prior[i])
+        gamma.star <- temp[1, ]
+        delta.star <- temp[2, ]
+      }
+      list(gamma.star=gamma.star, delta.star=delta.star)
+    }, BPPARAM = BPPARAM)
+    for (i in 1:n.batch) {
+      gamma.star[i,] <- results[[i]]$gamma.star
+      delta.star[i,] <- results[[i]]$delta.star
     }
-      
-    ## Find EB batch adjustments
-
-    gamma.star <- delta.star <- matrix(NA, nrow=n.batch, ncol=nrow(s.data))
-    if (par.prior) {
-        message("Finding parametric adjustments")
-        results <- bplapply(1:n.batch, function(i) {
-            if (mean.only) {
-                gamma.star <- postmean(gamma.hat[i,], gamma.bar[i], 1, 1, t2[i])
-                delta.star <- rep(1, nrow(s.data))
-            }
-            else {
-                temp <- it.sol(s.data[, batches[[i]]], gamma.hat[i, ],
-                               delta.hat[i, ], gamma.bar[i], t2[i], a.prior[i],
-                               b.prior[i])
-                gamma.star <- temp[1, ]
-                delta.star <- temp[2, ]
-            }
-            list(gamma.star=gamma.star, delta.star=delta.star)
-        }, BPPARAM = BPPARAM)
-        for (i in 1:n.batch) {
-            gamma.star[i,] <- results[[i]]$gamma.star
-            delta.star[i,] <- results[[i]]$delta.star
-        }
+  }   else {
+    message("Finding nonparametric adjustments")
+    results <- bplapply(1:n.batch, function(i) {
+      if (mean.only) {
+        delta.hat[i, ] = 1
+      }
+      temp <- int.eprior(as.matrix(s.data[, batches[[i]]]),
+                         gamma.hat[i, ], delta.hat[i, ])
+      list(gamma.star=temp[1,], delta.star=temp[2,])
+    }, BPPARAM = BPPARAM)
+    for (i in 1:n.batch) {
+      gamma.star[i,] <- results[[i]]$gamma.star
+      delta.star[i,] <- results[[i]]$delta.star
     }
-    else {
-        message("Finding nonparametric adjustments")
-        results <- bplapply(1:n.batch, function(i) {
-            if (mean.only) {
-                delta.hat[i, ] = 1
-            }
-            temp <- int.eprior(as.matrix(s.data[, batches[[i]]]),
-                               gamma.hat[i, ], delta.hat[i, ])
-            list(gamma.star=temp[1,], delta.star=temp[2,])
-        }, BPPARAM = BPPARAM)
-        for (i in 1:n.batch) {
-            gamma.star[i,] <- results[[i]]$gamma.star
-            delta.star[i,] <- results[[i]]$delta.star
-        }
+  }
+  
+  if(!is.null(ref.batch)){
+    gamma.star[ref,] <- 0  ## set reference batch mean equal to 0
+    delta.star[ref,] <- 1  ## set reference batch variance equal to 1
+  }
+  
+  ## Normalize the Data ###
+  message("Adjusting the Data\n")
+  
+  bayesdata <- s.data
+  
+  for (j in 1:n.batch){
+    i <- batches[[j]]
+    bayesdata[,i] <- (bayesdata[,i]-t(batch.design[i,]%*%gamma.star))/(sqrt(delta.star[j,])%*%t(rep(1,n.batches[j]))) # FIXME
+  }
+  
+  
+  bayesdata <- (bayesdata*(sqrt(var.pooled)%*%t(rep(1,n.array))))+stand.mean # FIXME
+  
+  ## Do not change ref batch at all in reference version
+  if(!is.null(ref.batch)){
+    bayesdata[, batches[[ref]]] <- dat[, batches[[ref]]]
+  }
+  
+  ## put genes with 0 variance in any batch back in data
+  if (length(zero.rows) > 0) {
+    dat.orig[keep.rows, ] <- bayesdata
+    bayesdata <- dat.orig   
+  }
+  
+  if(!is.null(dat_test)) {
+    
+    dat2 <- as.matrix(dat_test)
+    if( (nlevels(as.factor(batch_test)) > nlevels (batch)) | ( length(intersect(levels(as.factor(batch_test)), levels(batch))) <2 )) {
+      stop("batch_test should contain at least two batches in common with batch")
     }
     
-    if(!is.null(ref.batch)){
-        gamma.star[ref,] <- 0  ## set reference batch mean equal to 0
-        delta.star[ref,] <- 1  ## set reference batch variance equal to 1
+    batch2 <- factor(batch_test,levels=levels(batch))
+    
+    if (length(zero.rows) > 0) {
+      # keep a copy of the original data matrix and remove zero var rows
+      dat.orig2 <- dat2
+      dat2 <- dat2[keep.rows, ]
     }
+    batchmod2 <- model.matrix(~-1+batch2)  
+    ## A few other characteristics on the batches
+    batches2 <- list()
+    for (i in 1:n.batch)      batches2[[i]] <- which(batch2 == levels(batch2)[i])
+    n.batches2 <- sapply(batches2, length)
+    n.array2 <- sum(n.batches2)
+    ## combine batch variable and covariates
+    design2 <- cbind(batchmod2,mod)
+    design2 <- as.matrix(design2[,!check])
+    
+    stand.mean2 <- t(grand.mean) %*% t(rep(1,n.array2)) 
     
-    ## Normalize the Data ###
-    message("Adjusting the Data\n")
-  
-    bayesdata <- s.data
-    j <- 1
-    for (i in batches){
-        bayesdata[,i] <- (bayesdata[,i]-t(batch.design[i,]%*%gamma.star))/(sqrt(delta.star[j,])%*%t(rep(1,n.batches[j]))) # FIXME
-        j <- j+1
+    if(!is.null(design)){
+      tmp <- design2
+      tmp[,c(1:n.batch)] <- 0
+      stand.mean2 <- stand.mean2+t(tmp %*% B.hat)
+    }  
+    
+    s.data2 <- (dat2-stand.mean2)/(sqrt(var.pooled) %*% t(rep(1,n.array2)))
+    
+    ##Get regression batch effect parameters
+    batch.design2 <- design2[, 1:n.batch]
+    
+    bayesdata2 <- s.data2
+    
+    for (j in 1:n.batch){
+      i <- batches2[[j]]
+      bayesdata2[,i] <- (bayesdata2[,i]-t(batch.design2[i,]%*%gamma.star))/(sqrt(delta.star[j,])%*%t(rep(1,n.batches2[j]))) # FIXME
     }
     
-    bayesdata <- (bayesdata*(sqrt(var.pooled)%*%t(rep(1,n.array))))+stand.mean # FIXME
-  
+    bayesdata2 <- (bayesdata2*(sqrt(var.pooled)%*%t(rep(1,n.array2))))+stand.mean2 # FIXME
+    
     ## Do not change ref batch at all in reference version
     if(!is.null(ref.batch)){
-        bayesdata[, batches[[ref]]] <- dat[, batches[[ref]]]
+      bayesdata2[, batches2[[ref]]] <- dat2[, batches2[[ref]]]
     }
     
-    ## put genes with 0 variance in any batch back in data
     if (length(zero.rows) > 0) {
-      dat.orig[keep.rows, ] <- bayesdata
-      bayesdata <- dat.orig
+      dat.orig2[keep.rows, ] <- bayesdata2
+      bayesdata2 <- dat.orig2
     }
-    
+    return(list(tr_data=bayesdata,te_data=bayesdata2))
+  } else {
     return(bayesdata)
+  }
 }