require(pacman)
source('../R/utils.R')

imp_vars <- function(imp, xvars) {
  names <- row.names(imp)
  tmp <- rep(0, length(xvars))
  names(tmp) <- xvars
  for (n in names) {
    if (grepl(":", n)) {
      nt <- strsplit(n, ':')[[1]][1]
      stopifnot(nt %in% xvars) 
      tmp[nt] = tmp[nt] + imp[n, 1]
    } else {
      tmp[n] = tmp[n] + imp[n,1]
    }
  }
  tmp  
}

glmnet_imp <- function(x, y, inter_x=FALSE, ...) {
  p_load(caret)
  p_load(glmnet)
  
  xvars <- colnames(x)
  
  cvars <- list()
  j = 1
  for (i in seq_along(xvars)) {
    xv <- xvars[i]
    if (class(x[,xv]) == 'factor') {
        xl = levels(x[,xv])
        if (length(xl) == 1) next 
        cvars[[xv]] <- list(oname = xv, 
                            levels = xl,
                            N = length(xl) - 1)
        x[,xv] <- factor(x[,xv], levels = xl, labels = paste0(":",1:length(xl)))
        j = j+1
      }
  }
  data = data.frame(x, y)
  xtrain <- scale(model.matrix(y~-1 + ., data = data), TRUE, TRUE)
  
  rmi = c()
  for (i in 1:NCOL(xtrain)) {
    if (anyNA(xtrain[, i])) {
      rmi = c(rmi, i)
    }
  }
  if (!is.null(rmi)) {
    xtrain <- xtrain[, -rmi]
  }
  ytrain <- as.vector(y)
  
  family = ifelse(class(y) == 'factor', 'binomial', 'gaussian')
  
  cv.glmmod <- cv.glmnet(xtrain, ytrain, alpha=1, family = family)
  best.lambda <- cv.glmmod$lambda.min
  glmmod <- glmnet(xtrain, ytrain, 
                   alpha=1, lambda = best.lambda, family = family)
  imp = varImp(glmmod, lambda=best.lambda)
  imp = imp_vars(imp, xvars)
  
  for (xv in names(cvars)) {
    imp[xv] = imp[xv] / cvars[[xv]]$N
  }

  imp[abs(imp) < 1e-5] = 0  
  list(errRank = imp_measure(imp[xvars]), 
       imp = imp_scale(imp[xvars]),
       oimp = imp[xvars],
       mi = names(which.max(imp)))
}