/**
 * Copyright 2014-2026, XGBoost Contributors
 *
 * @brief General primitive for tree learning,
 *   Updating a collection of trees given the information.
 * \author Tianqi Chen
 */
#ifndef XGBOOST_TREE_UPDATER_H_
#define XGBOOST_TREE_UPDATER_H_

#include <dmlc/registry.h>
#include <xgboost/base.h>                // for Args, GradientPair
#include <xgboost/data.h>                // for DMatrix
#include <xgboost/gradient.h>            // for GradientContainer
#include <xgboost/host_device_vector.h>  // for HostDeviceVector
#include <xgboost/linalg.h>              // for VectorView
#include <xgboost/model.h>               // for Configurable
#include <xgboost/span.h>                // for Span
#include <xgboost/tree_model.h>          // for RegTree

#include <functional>  // for function
#include <string>      // for string
#include <vector>      // for vector

namespace xgboost {
namespace tree {
struct TrainParam;
}

class Json;
struct Context;
struct ObjInfo;

/**
 * @brief interface of tree update module, that performs update of a tree.
 */
class TreeUpdater : public Configurable {
 protected:
  Context const* ctx_ = nullptr;

 public:
  explicit TreeUpdater(const Context* ctx) : ctx_(ctx) {}
  ~TreeUpdater() override = default;
  /**
   * @brief Initialize the updater with given arguments.
   * @param args arguments to the objective function.
   */
  virtual void Configure(const Args& args) = 0;
  /**
   * @brief Whether this updater can be used for updating existing trees.
   *
   *  Some updaters are used for building new trees (like `hist`), while some others are
   *  used for modifying existing trees (like `prune`).  Return true if it can modify
   *  existing trees.
   */
  [[nodiscard]] virtual bool CanModifyTree() const { return false; }
  /**
   * @brief Whether the out_position in `Update` is valid. This determines whether adaptive
   *        tree can be used.
   */
  [[nodiscard]] virtual bool HasNodePosition() const { return false; }
  /**
   * @brief perform update to the tree models
   *
   * @param param  Hyper-parameter for constructing trees.
   * @param gpair  The gradient pair statistics of the data
   * @param p_fmat The data matrix passed to the updater.
   * @param out_position The leaf index for each row.  The index is negated if that row is
   *                     removed during sampling. So the 3th node is ~3.
   * @param out_trees references the trees to be updated, updater will change the content of trees
   *
   * @note All the trees in the vector are updated, with the same statistics, but maybe
   *       different random seeds, usually one tree is passed in at a time, there can be
   *       multiple trees when we train a random forest style model.
   */
  virtual void Update(tree::TrainParam const* param, GradientContainer* gpair, DMatrix* p_fmat,
                      common::Span<HostDeviceVector<bst_node_t>> out_position,
                      std::vector<RegTree*> const& out_trees) = 0;

  /**
   * @brief Determines whether updater has enough knowledge about a given dataset to
   *        quickly update prediction cache for the training data and performs the update
   *        if possible.
   *
   * @param p_fmat data matrix
   * @param out_preds prediction cache to be updated
   *
   * @return boolean indicating whether updater has capability to update the prediction
   *         cache. If true, the prediction cache will have been updated by the time this
   *         function returns.
   */
  virtual bool UpdatePredictionCache(DMatrix const* /*data*/,
                                     common::Span<HostDeviceVector<bst_node_t>> /*out_position*/,
                                     linalg::MatrixView<float> /*out_preds*/) {
    return false;
  }

  [[nodiscard]] virtual char const* Name() const = 0;

  /**
   * @brief Create a tree updater given name
   *
   * @param name Name of the tree updater.
   * @param ctx A global runtime parameter
   * @param task Infomation about the objective.
   */
  static TreeUpdater* Create(const std::string& name, Context const* ctx, ObjInfo const* task);
};

/**
 * @brief Registry entry for tree updater.
 */
struct TreeUpdaterReg
    : public dmlc::FunctionRegEntryBase<
          TreeUpdaterReg, std::function<TreeUpdater*(Context const* ctx, ObjInfo const* task)>> {};

/*!
 * \brief Macro to register tree updater.
 *
 * \code
 * // example of registering a objective ndcg@k
 * XGBOOST_REGISTER_TREE_UPDATER(ColMaker, "colmaker")
 * .describe("Column based tree maker.")
 * .set_body([]() {
 *     return new ColMaker<TStats>();
 *   });
 * \endcode
 */
#define XGBOOST_REGISTER_TREE_UPDATER(UniqueId, Name)                   \
  static DMLC_ATTRIBUTE_UNUSED ::xgboost::TreeUpdaterReg&               \
  __make_ ## TreeUpdaterReg ## _ ## UniqueId ## __ =                    \
      ::dmlc::Registry< ::xgboost::TreeUpdaterReg>::Get()->__REGISTER__(Name)

}  // namespace xgboost
#endif  // XGBOOST_TREE_UPDATER_H_
