library(jsonlite)
library(dplyr)
library(stringplus)
library(this.path)
setwd(this.path::this.dir())

# append function such that append(x) <- "y" is equivalent to x <- c(x, "y")
`a<-` <- function (x, value) c(x, value)

# read command line args
# args <- commandArgs(trailingOnly = TRUE)
# print(args)
# model_path <- args[1]
# output_path <- args[2]
# standalone <- args[3] %>% as.logical

model_path <- "data/blockshuffle_xgboost_model.json"
output_path <- "../../../src/io/xgboost_blockshuffle_model.h"
standalone <- FALSE
# output_path <- "xgboost_model_cpp_standalone.cpp"
# standalone <- TRUE

if(is.na(standalone)) stop("3rd argument must be TRUE or FALSE, whether the output is a standalone Rcpp file or not")

model <- jsonlite::read_json(model_path, simplify=FALSE)
ntrees <- model$learner$attributes$best_ntreelimit %>% as.numeric
trees <- model$learner$gradient_booster$model$trees
base_score <- model$learner$learner_model_param$base_score

# highest number of nodes in any tree
max_tree_size <- sapply(trees, function(x) x$left_children %>% length) %>% max

node_index_type <- "uint16_t"
null_node_index <- 65535
null_feature_index <- 255 # we only have 9 features

# write script call with args
OUT <- c("// auto-generated by convert_xgboost_json.R", 
         "// Rscript convert_xgboost_json.R %s %s %s" | c(model_path, output_path, standalone))

if(standalone) {
  a(OUT) <- c("#include <array>", "#include <memory>", "#include <vector>")
  a(OUT) <- c("#include <Rcpp.h>", "using namespace Rcpp;")
} else {
  a(OUT) <- c("#ifndef _QS2_XGBOOST_BLOCKSHUFFLE_MODEL_H", "#define _QS2_XGBOOST_BLOCKSHUFFLE_MODEL_H")
  a(OUT) <- c("#include <array>", "#include <memory>", "#include <vector>", '#include "io_common.h"')
}
  
a(OUT) <- '

namespace XgboostBlockshuffleModel {

struct XgTree {
  struct XgNode {
    XgNode * left_child;
    XgNode * right_child;
    uint8_t split_idx;
    double split_cond;
  };
  static constexpr uint8_t NULL_FEATURE_IDX = 255;
  static constexpr uint16_t NULL_NODE_INDEX = 65535;
  std::unique_ptr<XgNode[]> nodes;
  // constructor takes initializer lists
  XgTree(const std::vector<double> & split_conditions,
         const std::vector<uint8_t> & split_indices,
         const std::vector<uint16_t> & left_children,
         const std::vector<uint16_t> & right_children) :
  nodes(MAKE_UNIQUE_BLOCK_CUSTOM(XgNode, split_conditions.size())) {
    for(size_t i = 0; i < split_conditions.size(); ++i) {
      nodes[i].split_idx = split_indices[i];
      nodes[i].split_cond = split_conditions[i];
      nodes[i].left_child = left_children[i] == NULL_NODE_INDEX ? nullptr : &nodes[left_children[i]];
      nodes[i].right_child = right_children[i] == NULL_NODE_INDEX ? nullptr : &nodes[right_children[i]];
    }
  }
  template <typename T> double predict(const T & features) const {
    static_assert(std::is_same<typename T::value_type, double>::value, "T::value_type must be double");
    XgNode * current_node = &nodes[0];
    while(true) {
      if(current_node->split_idx == NULL_FEATURE_IDX) {
        return current_node->split_cond;
      }
      if(features[current_node->split_idx] < current_node->split_cond) {
        current_node = current_node->left_child;
      } else {
        current_node = current_node->right_child;
      }
    }
  }
};
'

nfeatures <- 0
a(OUT) <- "static constexpr float base_score = %s;" | base_score
a(OUT) <- "static const std::array<XgTree, %s> XgForest = {" | ntrees
for(i in 1:ntrees) {
    tree <- trees[[i]]
    left_children <- tree$left_children %>% unlist %>% ifelse(.==-1, null_node_index, .)
    right_children <- tree$right_children %>% unlist %>% ifelse(.==-1, null_node_index, .)
    split_conditions <- tree$split_conditions %>% unlist
    split_indices <- tree$split_indices %>% unlist %>% ifelse(left_children == null_node_index & right_children == null_node_index, null_feature_index, .)

    # the number of features is the max+1 out of all the split_indices
    nfeatures <- c(nfeatures, split_indices+1) %>% unique %>% setdiff(null_feature_index+1) %>% max

    len <- length(left_children)
    stopifnot(len == length(right_children))
    stopifnot(len == length(split_conditions))
    stopifnot(len == length(split_indices))

    # convert to strings
    left_children <- left_children %>% paste0(collapse = ",")
    right_children <- right_children %>% paste0(collapse = ",")
    split_conditions <- split_conditions %>% paste0(collapse = ",")
    split_indices <- split_indices %>% paste0(collapse = ",")
    # wrap in initializer list
    if(i == ntrees) { # no comma for last element
      line <- "XgTree({%s}, {%s}, {%s}, {%s})" | c(split_conditions, split_indices, left_children, right_children)
    } else {
      line <- "XgTree({%s}, {%s}, {%s}, {%s})," | c(split_conditions, split_indices, left_children, right_children)
    }
    a(OUT) <- line
}
a(OUT) <- "};" # end XgForest
a(OUT) <- ""

a(OUT) <- "
double predict_xgboost_impl(const std::array<double, %s> & features) {
  double result = base_score;
  for (auto & tree : XgForest) {
    result += tree.predict(features);
  }
  return result;
}

} // end namespace
" | nfeatures


if(standalone) {
a(OUT) <- "
// [[Rcpp::export(rng=false)]]
NumericVector predict_xgboost_cpp(DataFrame features) {"
  for(i in 1:nfeatures) {
    a(OUT) <- "  NumericVector x%s = features[%s];" | c(i-1, i-1)
  }
  a(OUT) <- "
  size_t nx = Rf_xlength(x0);
  NumericVector out(nx);
  for(size_t i=0; i<nx; ++i) {
    std::array<double, %s> f = {" | nfeatures
  for(i in 1:nfeatures) {
    if(i == nfeatures) {
      a(OUT) <- "      x%s[i]" | (i-1)
    } else {
      a(OUT) <- "      x%s[i]," | (i-1)
    }

  }
  a(OUT) <- "    };
  out[i] = XgboostBlockshuffleModel::predict_xgboost_impl(f);
  }
  return out;
}
"
} else {
  a(OUT) <- "
#endif // include guard
"
}

writeLines(text = OUT, con = output_path)
