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

suppressWarnings(suppressMessages(library(glmnet)))
library(methods)
############################ helper functions #######################


# optimize alpha parameter
optimize <- function(genotype, phenotype, alpha=seq(0,1,0.1), repet=7) {
  acc <- NULL
  indexAlpha <- 1
  for(a in alpha) {
    curAcc <- NULL
    # repeat nfolds time each analysis
    for(i in 1:repet) {
      # draw at random 1/3 of the training set for testing and thus choose alpha
      # note it is not a cross-validation
      n <- ceiling(nrow(genotype)/3)
      indexTest <- sample(1:nrow(genotype), size=n)
      # create training set and test set
      train <- genotype[-indexTest,]
      test <- genotype[indexTest,]
      phenoTrain <- phenotype[-indexTest]
      phenoTest <- phenotype[indexTest]
      # cv.glmnet allow to compute lambda at the current alpha
      cv <- cv.glmnet(x=as.matrix(train), y=phenoTrain, alpha=a)
      # create model
      model <- glmnet(x=as.matrix(train), y=phenoTrain, alpha=a, lambda = cv$lambda.1se)
      # predict test set
      pred <- predict(model, test, type = "response")
      # compute r2 for choosing the best alpha
      curAcc <- c(curAcc, r2(phenoTest, pred))
    }
    # add mean r2 for this value of lambda to the accuracy vector
    acc <- c(acc, mean(curAcc))
  }
  # choose best alpha
  names(acc) <- alpha
  return(as.numeric(names(acc)[which.max(acc)]))
}

# 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 ###########################

lasso <- function(genotype, phenotype, evaluation = T, outFile, folds, alpha=NULL) {
  # go for optimization
  if(is.null(alpha)) {
    alpha <- seq(0,1,0.1)
    alpha <- optimize(genotype=genotype, phenotype=phenotype, alpha = alpha)
  }
  # evaluation
  if(evaluation) {
    prediction <- NULL
    # do cross-validation
    for(i in 1:length(folds)) {
      # create training and test set
      train <- genotype[-folds[[i]],]
      test <- genotype[folds[[i]],]
      phenoTrain <- phenotype[-folds[[i]]]
      phenoTest <- phenotype[folds[[i]]]
      # cv.glmnet helps to compute the right lambda for a chosen alpha
      cv <- cv.glmnet(x=as.matrix(train), y=phenoTrain, alpha=alpha)
      # create model
      lasso.fit <- glmnet(x=as.matrix(train), y=phenoTrain, alpha=alpha, lambda = cv$lambda.1se)
      # predict value of the test set for further evaluation
      prediction <- c(prediction, list(predict(lasso.fit, test, type = "response")[,1]))
    }
    # save predicted value for test set of each fold for further evaluation
    saveRDS(prediction, file=paste(outFile,".rds", sep=""))
    # just create a model
  } else {
    # cv.glmnet helps to compute the right lambda for a chosen alpha
    cv <- cv.glmnet(x=genotype, y=phenotype, alpha=alpha)
    # create model
    model <- glmnet(x=genotype, y=phenotype, alpha=alpha, lambda=cv$lambda.1se)
    # save model
    saveRDS(model, file = paste(outFile, ".rds", sep = ""))
  }
}

############################ main #############################
# load argument
cmd <- commandArgs(T)
source(cmd[1])
# check if evaluation is required
evaluation <- F
if(as.integer(doEvaluation) == 1) {
  evaluation <- T
  con = file(folds)
  folds <- readLines(con = con, n = 1, ok=T)
  close(con)
  folds <- readRDS(folds)
}
# load classifier parameters
alpha <- as.numeric(alpha)
if(alpha < 0 | alpha > 1) {alpha <- NULL}
# 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 !
lasso(genotype = data.matrix(genotype), phenotype = phenotype,
               evaluation = evaluation, outFile = out, folds = folds, alpha = alpha)
# return path of the result file to galaxy
cat(paste(paste(out, ".rds", sep = ""), "\n", sep=""))