########################################################
#
# creation date : 05/01/16
# last modification : 02/09/16
# author : Dr Nicolas Beaume
#
########################################################


#log <- file(paste(getwd(), "log_rrBLUP.txt", sep="/"), open = "wt")
#sink(file = log, type="message")
library(rrBLUP)
############################ helper functions #######################


################################## main function ###########################

rrBLUP <- function(genotype, phenotype, outFile, evaluation=F, folds) {
  if(evaluation) {
    prediction <- NULL
    for(i in 1:length(folds)) {
      train <- genotype[-folds[[i]],]
      test <- genotype[folds[[i]],]
      phenoTrain <- phenotype[-folds[[i]]]
      phenoTest <- phenotype[folds[[i]]]
      model <-mixed.solve(phenoTrain, Z=genoTrain,K=NULL, SE=F,return.Hinv = F)
      pred <- as.matrix(genoTest) %*% as.matrix(model$u)
      pred <- pred[,1]+model$beta
      prediction <- c(prediction, list(pred))
    }
    saveRDS(prediction, file=paste(outFile,".rds", sep=""))
    # just create a model
  } else {
    model <-mixed.solve(phenotype, Z=genotype,K=NULL, SE=F,return.Hinv = F)
    saveRDS(model, file = paste(outFile, ".rds", sep = ""))
  }
}


############################ main #############################
# running from terminal (supposing the OghmaGalaxy/bin directory is in your path) :
# rrBLUP.sh -i path_to_genotype -p phenotype_file -f folds -e -f fold_file -o out_file 
## -i : path to the file that contains the genotypes, must be a .rda file (as outputed by loadGenotype).
# please note that the table must be called "encoded" when your datafile is saved into .rda (automatic if encoded.R was used)

## -p : file that contains the phenotype must be a .rda file (as outputed by loadGenotype.R).
# please note that the table must be called "phenotype" when your datafile is saved into .rda (automatic if loadGenotype.R was used)

## -e : do evaluation instead of producing a model

## -f : index of the folds to which belong each individual
# please note that the list must be called "folds" (automatic if folds.R was used)

## -o : path to the output folder and generic name of the analysis. The extension related to each type of
# files are automatically added

cmd <- commandArgs(T)
source(cmd[1])
# load genotype and phenotype
con = file(genotype)
genotype <- readLines(con = con, n = 1, ok=T)
close(con)
genotype <- read.table(genotype, sep="\t", h=T)
# phenotype is written as a table (in columns) but it must be sent as a vector for mixed.solve
phenotype <- read.table(phenotype, sep="\t", h=T)[,1]
# run !
rrBLUP(genotype = genotype, phenotype = phenotype, outFile = out)
cat(paste(paste(out, ".rds", sep = ""), "\n", sep=""))