# Forward Recursive Feature Selection (Version 1.1)
# Developed by Songchao Chen (chensongchao@zju.edu.cn)
# Date: 24th August 2022
# Compared to Version 1.0, this version uses ranger package to speed the computation for large dataset

FRFS <- function(data, var.res, method="rf", early.stop=NULL){
  # data is the data.frame that contains response variable (var.res, or Y) and explanatory variables (var.exp, or X)
  # var.res is the response variable to be predicted (or Y)
  # method indicates the model to be used, the default model is random forest. The current version only use random forest.
  # early.stop indicates whether stop the algorithm when the reach a local maximum of accuracy (in RMSE). It is recommended when the number of variables is greater than 50 to speed up the calculation.
  require(caret)
  require(ranger)
  
  # check whether use random forest
  if(method!="rf"){
    message("Error! The current version only supports random forest model.")
  }
  
  # get the explanatory variables
  col.var.res <- which(colnames(data)==var.res)
  var.exp <- colnames(data)[-col.var.res]
  
  # determine the most important explanatory variables 
  set.seed(666)
  rf.temp <- ranger(x=data[,-col.var.res], y=data[,col.var.res], num.trees=500, importance="permutation")
  vip <- as.data.frame(rf.temp$variable.importance)
  var.choose <- rownames(vip)[which(vip[,1]==max(vip[,1]))]
  var.left <- var.exp[-which(var.exp %in% var.choose)]
  
  # determine the optimal variables
  set.seed(666)
  fitControl <- trainControl(method = "cv",number = 5) # 5-fold cross-validation
  Grid <- expand.grid(mtry=seq(1,length(var.choose),1), min.node.size=5, splitrule="variance")
  formula <- as.formula(paste0(var.res, "~", paste(var.choose, collapse="+")))
  rf.temp <- train(formula, data=data, method="ranger", trControl=fitControl, tuneGrid=Grid)
  RMSE.all <- data.frame(No.var=1:length(var.exp),RMSE=NA)
  RMSE.all[length(var.choose),2] <- rf.temp$results$RMSE
  
  for(iter in 2:length(var.exp)){
    RMSE.temp <- data.frame(Var=var.left, RMSE=NA)
    for(i in 1:length(var.left)){
      var.test <- c(var.choose, var.left[i])
      set.seed(666)
      fitControl <- trainControl(method = "cv",number = 5) # 5-fold cross-validation
      Grid <- expand.grid(mtry=seq(1,length(var.test),1), min.node.size=5, splitrule="variance")
      formula <- as.formula(paste0(var.res, "~", paste(var.test, collapse="+")))
      rf.temp <- train(formula, data=data, method="ranger", trControl=fitControl, tuneGrid=Grid)
      RMSE.temp[i,2] <- min(rf.temp$results$RMSE)
    }

    if(min(RMSE.temp[,2]) < min(RMSE.all[,2], na.rm=TRUE)){
      var.choose <- c(var.choose, RMSE.temp[which(RMSE.temp[,2]==min(RMSE.temp[,2])),1])
      var.left <- var.exp[-which(var.exp %in% var.choose)]
      RMSE.all[length(var.choose),2] <- min(RMSE.temp[,2])
    }else{
      if(early.stop == TRUE){ # Early stop
        RMSE.all[length(var.choose)+1,2] <- min(RMSE.temp[,2])
        break
      }else{
        var.choose <- c(var.choose, RMSE.temp[which(RMSE.temp[,2]==min(RMSE.temp[,2])),1])
        var.left <- var.exp[-which(var.exp %in% var.choose)]
        RMSE.all[length(var.choose),2] <- min(RMSE.temp[,2])
      }
    }
  }
  
  # export the final selected variables
  if(early.stop==FALSE){
    var.final <- var.choose[1:which(RMSE.all[,2]==min(RMSE.all[,2]))]
  }else{
    var.final <- var.choose
  }
  
  return(list(var.order=var.choose, var.final=var.final, RMSE=RMSE.all))
}