require(pacman)

p_load(caret)

decide_cat <- function(cat_x) {
  NL = 0:(NCOL(cat_x) - 1)
  dat = apply(cat_x, 1, FUN = function(x) {NL[which.max(x)]})
  factor(dat)
}

mean_impute <- function(x) {
  mu = mean(x, na.rm=TRUE)
  x[is.na(x)] = mu
  return(x)
}

median_impute <- function(x) {
  mu = median(x, na.rm = TRUE)
  x[is.na(x)] = mu
  return(x)
}

mode_impute <- function(x) {
  getmode <- function(v) {
    uniqv <- unique(v)
    uniqv[which.max(tabulate(match(v, uniqv)))]
  }
  mu <-getmode(x[!is.na(x)])
  x[is.na(x)] = mu
  return(x)
}

naive_impute <- function(x) {
  P = NCOL(x)
  for (i in 1:P) {
    if (is.numeric(x[,i])) {
      x[,i] = mean_impute(x[,i])
    } else {
      x[,i] = mode_impute(x[,i])
    }
  }
  return(x)
}

miss_important <- function(x, y, method, miss = 'naive', ...) {
  if (miss == 'naive') {
    xImp = naive_impute(x)
    return(method(xImp, y, ...))
  }
  suppressWarnings({imp = multiple_impute(x)})
  ans = matrix(NA, 5, NCOL(x))
  for (i in seq_along(imp)) {
    xImp = imp[[i]]
    ans[i,] = method(xImp, y, ...)$oimp
  }
  cmu = colMeans(ans)
  errRank = imp_measure(cmu)
  
  return(list(errRank = errRank, 
              imp = imp_scale(cmu),
              oimp = cmu))
}

boot_correction <- function(x, y, method, B = 5, ...) {
  P = NCOL(x)
  N = NROW(x)
  res = matrix(0, B, P)
  ori = method(x, y, ...)$oimp
  for (i in 1:B) {
    idx = sample.int(n = N, size = N, replace = TRUE)
    tmp = method(x, y[idx])
    res[i,] = tmp$oimp
  }
  cmu = colMeans(res)
  csd = apply(res, 2, sd)
  #idx = which(csd < 1e-5)
  #cmu[idx] = 0
  #csd[idx] = 1
  res = (ori - cmu) / csd
  names(res) = names(ori)
  imp_scale(res)
}


caret_impute <- function(x) {
  xvars = colnames(x)
  dummies = dummyVars(~., data = x)
  dumx = predict(dummies, x)
  preProcValues <- preProcess(dumx, method = c("bagImpute"))
  impDumx = predict(preProcValues, dumx)

  vars = colnames(impDumx)
  for (xv in xvars) {
    if (!anyNA(x[[xv]])) {next}
    if (class(x[[xv]]) %in% c("numeric", "integer")) {
      x[,xv] = impDumx[, xv]
      next
    } else {
    xCv = vars[grepl(xv, vars)]
    ord = order(xCv)
    xCv = xCv[ord]
    xCd = impDumx[, xCv]
    x[xv] = decide_cat(xCd)
    }
  }
  return(x)
}

