########################################################
#
# creation date : 08/01/16
# last modification : 23/06/16
# author : Dr Nicolas Beaume
# owner : IRRI
#
########################################################

# compute correlation
predictionCorrelation <- function(prediction, obs) {
  return(cor(prediction, obs, method = "pearson"))
}

# compute r2 by computing the classic formula
# compare the sum of square difference from target to prediciton
# to the sum of square difference from target to the mean of the target
r2 <- function(target, prediction) {
  sst <- sum((target-mean(target))^2)
  ssr <- sum((target-prediction)^2)
  return(1-ssr/sst)
}

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

evaluatePredictions <- function(data, prediction=NULL, traitValue=NULL, doR2=F, 
                                folds=NULL, classes=NULL, doCor=F) {
  eval <- NULL
  # case of r2
  if(doR2) {
    eval <- c(eval, list(r2=NULL))
  }
  if(doCor) {
    eval <- c(eval, list(cor=NULL))
  }
  for (i in 1:length(folds)) {
    train <- unlist(folds[-i])
    test <- folds[[i]]
    if(doR2) {
      if(!is.null(traitValue)) {
        eval$r2 <- c(eval$r2, r2(traitValue[test], prediction[[i]]))
      } else {
        eval$r2 <- c(eval$r2, NA)
      }
    }
    # case of correlation
    if(doCor) {
      if(!is.null(traitValue)) {
        eval$cor <- c(eval$cor, cor(traitValue[test], prediction[[i]]))
      } else {
        eval$cor <- c(eval$cor, NA)
      }
    }
  }
  # print output
  print(eval)
  if(doR2) {
    print(paste("mean r2 : ", mean(eval$r2), "standard deviation :", sd(eval$r2)))
  }
  if(doCor) {
    print(paste("mean correlation : ", mean(eval$cor), "standard deviation :", sd(eval$cor)))
  }
}
############################ main #############################

# load arguments
cmd <- commandArgs(trailingOnly = T)
source(cmd[1])
# set which evaluation are used
if(as.integer(doR2) == 1) {
  doR2 <- T
}
if(as.integer(doCor) == 1) {
  doCor <- T
}
# load genotype & phenotype
con = file(genotype)
genotype <- readLines(con = con, n = 1, ok=T)
close(con)
genotype <- read.table(genotype, sep="\t",h=T)
phenotype <- read.table(phenotype, sep="\t",h=T)[,1]
# load prediction
con = file(prediction)
prediction <- readLines(con = con, n = 1, ok=T)
close(con)
prediction <- readRDS(prediction)
# load folds
con = file(folds)
folds <- readLines(con = con, n = 1, ok=T)
close(con)
folds <- readRDS(folds)
# evaluation
evaluatePredictions(data=genotype, prediction=prediction, traitValue=phenotype, doR2=doR2, 
                    folds=folds, doCor=doCor)