// Train a decision tree regressor on random numeric data and make predictions.

// All data and responses are uniform random; this uses 10 dimensional data.
// Replace with a data::Load() call or similar for a real application.
arma::mat dataset(10, 1000, arma::fill::randu); // 1000 points.
arma::rowvec responses = arma::randn<arma::rowvec>(1000);
arma::mat testDataset(10, 500, arma::fill::randu); // 500 test points.

mlpack::DecisionTreeRegressor tree;     // Step 1: create tree.
tree.Train(dataset, responses);         // Step 2: train model.
arma::rowvec predictions;
tree.Predict(testDataset, predictions); // Step 3: use model to predict.

// Print some information about the test predictions.
std::cout << arma::accu(predictions > 0.7) << " test points predicted to have"
    << " responses greater than 0.7." << std::endl;
std::cout << arma::accu(predictions < 0) << " test points predicted to have "
    << "negative responses." << std::endl;

// Load a categorical dataset.
arma::mat data;
mlpack::data::DatasetInfo info;
// See https://datasets.mlpack.org/telecom_churn.arff.
mlpack::data::Load("telecom_churn.arff", data, info, true);

arma::rowvec responses;
// See https://datasets.mlpack.org/telecom_churn.responses.csv.
mlpack::data::Load("telecom_churn.responses.csv", responses, true);

// Split data into training set (80%) and test set (20%).
arma::mat trainData, testData;
arma::rowvec trainResponses, testResponses;
mlpack::data::Split(data, responses, trainData, testData, trainResponses,
    testResponses, 0.2);

// Create the tree.
mlpack::DecisionTreeRegressor tree;
// Train on the given dataset, specifying a minimum gain of 1e-6 and keeping the
// default minimum leaf size.
const double mse = tree.Train(trainData, info, trainResponses,
    10 /* minimum leaf size */, 1e-6 /* minimum gain */);
// Print the MSE of the trained tree.
std::cout << "MSE of trained tree is " << mse << "." << std::endl;

// Compute prediction on the first test point.
const double firstPrediction = tree.Predict(testData.col(0));
std::cout << "Predicted value for first test point is " << firstPrediction
    << "." << std::endl;

// Compute predictions on test data.
arma::rowvec testPredictions;
tree.Predict(testData, testPredictions);

// Compute the average error on the test set.
const double testAverageError = arma::mean(testResponses - testPredictions);
std::cout << "Average error on test set: " << testAverageError << "."
    << std::endl;

mlpack::DecisionTreeRegressor tree;
// This call assumes a tree called "tree" has already been saved to `tree.bin`
// with `data::Save()`.
mlpack::data::Load("tree.bin", "tree", tree, true);

std::cout << "Information about the DecisionTreeRegressor in `tree.bin`:"
    << std::endl;
std::cout << " * The root node has " << tree.NumChildren() << " children."
    << std::endl;
std::cout << " * The tree has " << tree.NumLeaves() << " leaves." << std::endl;
if (tree.NumChildren() > 0)
{
  for (size_t i = 0; i < tree.NumChildren(); ++i)
  {
    std::cout << " * Child " << i << " of the root has "
        << tree.Child(i).NumLeaves() << " leaves in its subtree." << std::endl;
  }
}

// 1000 random points in 10 dimensions.
arma::fmat dataset(10, 1000, arma::fill::randu);
// Random responses for each point, with a normal distribution.
arma::frowvec responses = arma::randn<arma::frowvec>(1000);

// Train in the constructor.
mlpack::DecisionTreeRegressor tree(dataset, responses, 5);

// Create test data (500 points).
arma::fmat testDataset(10, 500, arma::fill::randu);
arma::frowvec predictions;
tree.Predict(testDataset, predictions);
// Now `predictions` holds predictions for the test dataset.

// Print some information about the test predictions.
std::cout << arma::accu(predictions > 1) << " test points predicted to have "
    << "value greater than 1." << std::endl;

DecisionTreeRegressor<FitnessFunction,
                      NumericSplitType,
                      CategoricalSplitType,
                      DimensionSelectionType,
                      NoRecursion>

// You can use this as a starting point for implementation.
class CustomFitnessFunction
{
  // Compute the gain for the given vector of values, where `values[i]` has an
  // associated instance weight `weights[i]`.
  //
  // `RowType` and `WeightVecType` will be vector types following the Armadillo
  // API.  If `UseWeights` is `false`, then the `weights` vector should be
  // ignored (e.g. the responses are not weighted).
  //
  // In the version with `begin` and `end` parameters, only the subset between
  // `labels[begin]` and `labels[end]` (inclusive) should be considered.
  template<bool UseWeights, typename RowType, typename WeightVecType>
  double Evaluate(const RowType& labels,
                  const WeightVecType& weights);

  template<bool UseWeights, typename RowType, typename WeightVecType>
  double Evaluate(const RowType& labels,
                  const WeightVecType& weights,
                  const size_t begin,
                  const size_t end);

  // Return the output value for prediction for a leaf node whose training
  // values are made up of the values in the vector `responses` (optionally with
  // associated instance weights `weights`).
  //
  // `ResponsesType` and `WeightsType` will be vector types following the
  // Armadillo API.  If `UseWeights` is `false`, then the `weights` vector
  // should be ignored (e.g. the responses are not weighted).
  template<bool UseWeights, typename ResponsesType, typename WeightsType>
  double OutputLeafValue(const ResponsesType& responses,
                         const WeightsType& weights);
};

class CustomNumericSplit
{
 public:
  // If a split with better resulting gain than `bestGain` is found, then
  // information about the new, better split should be stored in `splitInfo` and
  // `aux`.  Specifically, a split is better than `bestGain` if the sum of the
  // gains that the children will have (call this `sumChildrenGains`) is
  // sufficiently better than the gain of the unsplit node (call this
  // `unsplitGain`):
  //
  //    split if `sumChildrenGains - unsplitGain > bestGain`, and
  //             `sumChildrenGains - unsplitGain > minGainSplit`, and
  //             each child will have at least `minLeafSize` points
  //
  // The new best split value should be returned (or anything greater than or
  // equal to `bestGain` if no better split is found).
  //
  // If a new best split is found, then `splitInfo` and `aux` should be
  // populated with the information that will be needed for
  // `CalculateDirection()` to successfully choose the child for a given point.
  // `splitInfo` should be set to a vector of length 1.  The format of `aux` is
  // arbitrary and is detailed more below.
  //
  // If `UseWeights` is false, the vector `weights` should be ignored.
  // Otherwise, they are instance weighs for each value in `data` (one dimension
  // of the input data).
  template<bool UseWeights, typename VecType, typename ResponsesType,
      typename WeightVecType>
  double SplitIfBetter(const double bestGain,
                       const VecType& data,
                       const ResponsesType& responses,
                       const WeightVecType& weights,
                       const size_t minLeafSize,
                       const double minGainSplit,
                       arma::vec& splitInfo,
                       AuxiliarySplitInfo& aux,
                       FitnessFunction& function);

  // Return the number of children for a given split (stored as the single
  // element from `splitInfo` and auxiliary data `aux` in `SplitIfBetter()`).
  size_t NumChildren(const double& splitInfo,
                     const AuxiliarySplitInfo& aux);

  // Given a point with value `point`, and split information `splitInfo` and
  // `aux`, return the index of the child that corresponds to the point.  So,
  // e.g., if the split type was a binary split on the value `splitInfo`, you
  // might return `0` if `point < splitInfo`, and `1` otherwise.
  template<typename ElemType>
  static size_t CalculateDirection(
      const ElemType& point,
      const double& splitInfo,
      const AuxiliarySplitInfo& /* aux */);

  // This class can hold any extra data that is necessary to encode a split.  It
  // should only be non-empty if a single `double` value cannot be used to hold
  // the information corresponding to a split.
  class AuxiliarySplitInfo { };
};

class CustomCategoricalSplit
{
 public:
  // If a split with better resulting gain than `bestGain` is found, then
  // information about the new, better split should be stored in `splitInfo` and
  // `aux`.  Specifically, a split is better than `bestGain` if the sum of the
  // gains that the children will have (call this `sumChildrenGains`) is
  // sufficiently better than the gain of the unsplit node (call this
  // `unsplitGain`):
  //
  //    split if `sumChildrenGains - unsplitGain > bestGain`, and
  //             `sumChildrenGains - unsplitGain > minGainSplit`, and
  //             each child will have at least `minLeafSize` points
  //
  // The new best split value should be returned (or anything greater than or
  // equal to `bestGain` if no better split is found).
  //
  // If a new best split is found, then `splitInfo` and `aux` should be
  // populated with the information that will be needed for
  // `CalculateDirection()` to successfully choose the child for a given point.
  // `splitInfo` should be set to a vector of length 1.  The format of `aux` is
  // arbitrary and is detailed more below.
  //
  // If `UseWeights` is false, the vector `weights` should be ignored.
  // Otherwise, they are instance weighs for each value in `data` (one
  // categorical dimension of the input data, which takes values between `0` and
  // `numCategories - 1`).
  template<bool UseWeights, typename VecType, typename ResponsesType,
           typename WeightVecType>
  static double SplitIfBetter(
      const double bestGain,
      const VecType& data,
      const size_t numCategories,
      const ResponsesType& labels,
      const WeightVecType& weights,
      const size_t minLeafSize,
      const double minGainSplit,
      arma::vec& splitInfo,
      AuxiliarySplitInfo& aux,
      FitnessFunction& fitnessFunction);

  // Return the number of children for a given split (stored as the single
  // element from `splitInfo` and auxiliary data `aux` in `SplitIfBetter()`).
  size_t NumChildren(const double& splitInfo,
                     const AuxiliarySplitInfo& aux);

  // Given a point with (categorical) value `point`, and split information
  // `splitInfo` and `aux`, return the index of the child that corresponds to
  // the point.  So, e.g., for `AllCategoricalSplit`, which splits a categorical
  // dimension into one child for each category, this simply returns `point`.
  template<typename ElemType>
  static size_t CalculateDirection(
      const ElemType& point,
      const double& splitInfo,
      const AuxiliarySplitInfo& /* aux */);

  // This class can hold any extra data that is necessary to encode a split.  It
  // should only be non-empty if a single `double` value cannot be used to hold
  // the information corresponding to a split.
  class AuxiliarySplitInfo { };
};

class CustomDimensionSelect
{
 public:
  // Get the first dimension to try.
  // This should return a value between `0` and `data.n_rows`.
  size_t Begin();

  // Get the next dimension to try.  Note that internal state can be used to
  // track which candidate dimension is currently being looked at.
  // This should return a value between `0` and `data.n_rows`.
  size_t Next();

  // Get a value indicating that all dimensions have been tried.
  size_t End() const;

  // The usage pattern of `DimensionSelectionType` by `DecisionTree` is as
  // follows, assuming that `dim` is an instantiated `DimensionSelectionType`
  // object:
  //
  // for (size_t dim = dim.Begin(); dim != dim.End(); dim = dim.Next())
  // {
  //   // ... try to split on dimension `dim` ...
  // }
};