######################################################################################
#                                                                                    #
# Simulation function for creating datasets and calculating VIs                      #
#                                                                                    #
# Input: - Number of observation in both classes (nobs1, nobs2)                      #
#                                                                                    #
#       (- number of noise predictors 'nnoise')                                      #
#       (- effects strengths 'effect1', 'effect2', 'effect3')                        #
#       (- number of iterations 'niter')                                             #
#       (- parameters controlling tree pruning)                                      #
#                                                                                    #
# Output: list of VIs using ER-based and AUC-based permutation VIM                   # 
#         1st list element: (npred x niter)–matrix of VI values using ER-based VIM   #
#         2nd list element: (npred x niter)–matrix of VI values using AUC-based VIM  #
#                                                                                    #
######################################################################################

require("party")

calc_VIMs <- function(nnoise = 50, nobs1 = 10, nobs2 = 10,
                      effect1 = 0.5, effect2 = 0.75, effect3 = 1,
                      minbuck = 0, minsp = 0, mincrit = 0, niter = 100){

  data <- vector("list", niter)
  res  <- vector("list", niter)

  ##################################################################
  # Loop for simulating niter datasets                             #
  ##################################################################

  for (i in 1 : niter){

    set.seed(i)   # set seed for reproducibility

    # Create predictors (5 of each effect strength and nnoise noise pred.)
    
    noise <- matrix(rnorm(nnoise * (nobs1 + nobs2)), ncol = nnoise)
    large.effects    <- rbind(matrix(rnorm(nobs1*5, mean = effect3), ncol = 5),
                              matrix(rnorm(nobs2*5, mean = 0      ), ncol = 5))
    moderate.effects <- rbind(matrix(rnorm(nobs1*5, mean = effect2), ncol = 5),
                              matrix(rnorm(nobs2*5, mean = 0      ), ncol = 5))
    small.effects    <- rbind(matrix(rnorm(nobs1*5, mean = effect1), ncol = 5),
                              matrix(rnorm(nobs2*5, mean = 0      ), ncol = 5))

    data[[i]] <- data.frame(y = as.factor(c(rep(1, nobs1), rep(0, nobs2))),
                            cbind(large.effects, moderate.effects, small.effects,
                                  noise))
  }

  ##################################################################
  # Function for calculating VIs for a given dataset               #
  ##################################################################

  get_VIMs <- function(xdata){

    # Grow random forest
    
    forest <- cforest(y ~ ., data = xdata, control =
                   cforest_control(teststat = "quad",
                   testtype     = "Univ",
                   minbucket    = minbuck,
                   minsplit     = minsp,
                   mincriterion = mincrit,
                   replace      = FALSE,
                   ntree        = 1000))

    # Calculate VIs
    
    varimp.ER  <- varimp(forest)
    varimp.AUC <- varimpAUC(forest)

    return(list(varimp.ER = varimp.ER, varimp.AUC = varimp.AUC))
    
  }

  ##################################################################
  # Call function get_VIMs() for each dataset                      #
  ##################################################################

  res <- lapply(data, get_VIMs)

  # Save VIs in matrices
  
  AUC_VI <- data.frame(res)[, grep("varimp.AUC", colnames(data.frame(res)))]
  ER_VI  <- data.frame(res)[, grep("varimp.ER",  colnames(data.frame(res)))]

  # Build list object containing VI values
  
  res    <- list(ER_VI = ER_VI, AUC_VI = AUC_VI)
  return(res)  # return list
  
}

