#ifndef XGBOOST_TREE_MODEL_H_
#define XGBOOST_TREE_MODEL_H_
/*!
 * \file model.h
 * \brief model structure for tree
 * \author Tianqi Chen
 */
#include <string>
#include <cstring>
#include <sstream>
#include <limits>
#include <algorithm>
#include <vector>
#include <cmath>
#include "../utils/io.h"
#include "../utils/fmap.h"
#include "../utils/utils.h"

namespace xgboost {
namespace tree {
/*!
 * \brief template class of TreeModel 
 * \tparam TSplitCond data type to indicate split condition
 * \tparam TNodeStat auxiliary statistics of node to help tree building
 */
template<typename TSplitCond, typename TNodeStat>
class TreeModel {
 public:
  /*! \brief data type to indicate split condition */
  typedef TNodeStat  NodeStat;
  /*! \brief auxiliary statistics of node to help tree building */
  typedef TSplitCond SplitCond;
  /*! \brief parameters of the tree */
  struct Param{
    /*! \brief number of start root */
    int num_roots;
    /*! \brief total number of nodes */
    int num_nodes;
    /*!\brief number of deleted nodes */
    int num_deleted;
    /*! \brief maximum depth, this is a statistics of the tree */
    int max_depth;
    /*! \brief  number of features used for tree construction */
    int num_feature;
    /*! 
     * \brief leaf vector size, used for vector tree
     * used to store more than one dimensional information in tree
     */
    int size_leaf_vector;
    /*! \brief reserved part */
    int reserved[31];
    /*! \brief constructor */
    Param(void) {
      max_depth = 0;
      size_leaf_vector = 0;
      std::memset(reserved, 0, sizeof(reserved));
    }
    /*! 
     * \brief set parameters from outside 
     * \param name name of the parameter
     * \param val  value of the parameter
     */
    inline void SetParam(const char *name, const char *val) {
      using namespace std;
      if (!strcmp("num_roots", name)) num_roots = atoi(val);
      if (!strcmp("num_feature", name)) num_feature = atoi(val);
      if (!strcmp("size_leaf_vector", name)) size_leaf_vector = atoi(val);
    }
  };
  /*! \brief tree node */
  class Node{
   public:
    /*! \brief index of left child */
    inline int cleft(void) const {
      return this->cleft_;
    }
    /*! \brief index of right child */
    inline int cright(void) const {
      return this->cright_;
    }
    /*! \brief index of default child when feature is missing */
    inline int cdefault(void) const {
      return this->default_left() ? this->cleft() : this->cright();
    }
    /*! \brief feature index of split condition */
    inline unsigned split_index(void) const {
      return sindex_ & ((1U << 31) - 1U);
    }
    /*! \brief when feature is unknown, whether goes to left child */
    inline bool default_left(void) const {
      return (sindex_ >> 31) != 0;
    }
    /*! \brief whether current node is leaf node */
    inline bool is_leaf(void) const {
      return cleft_ == -1;
    }
    /*! \brief get leaf value of leaf node */
    inline float leaf_value(void) const {
      return (this->info_).leaf_value;
    }
    /*! \brief get split condition of the node */
    inline TSplitCond split_cond(void) const {
      return (this->info_).split_cond;
    }
    /*! \brief get parent of the node */
    inline int parent(void) const {
      return parent_ & ((1U << 31) - 1);
    }
    /*! \brief whether current node is left child */
    inline bool is_left_child(void) const {
      return (parent_ & (1U << 31)) != 0;
    }
    /*! \brief whether current node is root */
    inline bool is_root(void) const {
      return parent_ == -1;
    }
    /*! 
     * \brief set the right child 
     * \param nide node id to right child
     */
    inline void set_right_child(int nid) {
      this->cright_ = nid;
    }
    /*! 
     * \brief set split condition of current node 
     * \param split_index feature index to split
     * \param split_cond  split condition
     * \param default_left the default direction when feature is unknown
     */
    inline void set_split(unsigned split_index, TSplitCond split_cond,
                          bool default_left = false) {
      if (default_left) split_index |= (1U << 31);
      this->sindex_ = split_index;
      (this->info_).split_cond = split_cond;
    }
    /*! 
     * \brief set the leaf value of the node
     * \param value leaf value
     * \param right right index, could be used to store 
     *        additional information
     */
    inline void set_leaf(float value, int right = -1) {
      (this->info_).leaf_value = value;
      this->cleft_ = -1;
      this->cright_ = right;
    }

   private:
    friend class TreeModel<TSplitCond, TNodeStat>;
    /*! 
     * \brief in leaf node, we have weights, in non-leaf nodes, 
     *        we have split condition 
     */
    union Info{
      float leaf_value;
      TSplitCond split_cond;
    };
    // pointer to parent, highest bit is used to
    // indicate whether it's a left child or not
    int parent_;
    // pointer to left, right
    int cleft_, cright_;
    // split feature index, left split or right split depends on the highest bit
    unsigned sindex_;
    // extra info
    Info info_;
    // set parent
    inline void set_parent(int pidx, bool is_left_child = true) {
      if (is_left_child) pidx |= (1U << 31);
      this->parent_ = pidx;
    }
  };

 protected:
  // vector of nodes
  std::vector<Node> nodes;
  // free node space, used during training process
  std::vector<int>  deleted_nodes;
  // stats of nodes
  std::vector<TNodeStat> stats;
  // leaf vector, that is used to store additional information
  std::vector<bst_float> leaf_vector;
  // allocate a new node,
  // !!!!!! NOTE: may cause BUG here, nodes.resize
  inline int AllocNode(void) {
    if (param.num_deleted != 0) {
      int nd = deleted_nodes.back();
      deleted_nodes.pop_back();
      --param.num_deleted;
      return nd;
    }
    int nd = param.num_nodes++;
    utils::Check(param.num_nodes < std::numeric_limits<int>::max(),
                 "number of nodes in the tree exceed 2^31");
    nodes.resize(param.num_nodes);
    stats.resize(param.num_nodes);
    leaf_vector.resize(param.num_nodes * param.size_leaf_vector); 
    return nd;
  }
  // delete a tree node
  inline void DeleteNode(int nid) {
    utils::Assert(nid >= param.num_roots, "can not delete root");
    deleted_nodes.push_back(nid);
    nodes[nid].set_parent(-1);
    ++param.num_deleted;
  }

 public:
  /*! 
   * \brief change a non leaf node to a leaf node, delete its children
   * \param rid node id of the node
   * \param new leaf value
   */
  inline void ChangeToLeaf(int rid, float value) {
    utils::Assert(nodes[nodes[rid].cleft() ].is_leaf(),
                  "can not delete a non termial child");
    utils::Assert(nodes[nodes[rid].cright()].is_leaf(),
                  "can not delete a non termial child");
    this->DeleteNode(nodes[rid].cleft());
    this->DeleteNode(nodes[rid].cright());
    nodes[rid].set_leaf(value);
  }
  /*! 
   * \brief collapse a non leaf node to a leaf node, delete its children
   * \param rid node id of the node
   * \param new leaf value
   */
  inline void CollapseToLeaf(int rid, float value) {
    if (nodes[rid].is_leaf()) return;
    if (!nodes[nodes[rid].cleft() ].is_leaf()) {
      CollapseToLeaf(nodes[rid].cleft(), 0.0f);
    }
    if (!nodes[nodes[rid].cright() ].is_leaf()) {
      CollapseToLeaf(nodes[rid].cright(), 0.0f);
    }
    this->ChangeToLeaf(rid, value);
  }

 public:
  /*! \brief model parameter */
  Param param;
  /*! \brief constructor */
  TreeModel(void) {
    param.num_nodes = 1;
    param.num_roots = 1;
    param.num_deleted = 0;
    nodes.resize(1);
  }
  /*! \brief get node given nid */
  inline Node &operator[](int nid) {
    return nodes[nid];
  }
  /*! \brief get node given nid */
  inline const Node &operator[](int nid) const {
    return nodes[nid];
  }
  /*! \brief get node statistics given nid */
  inline NodeStat &stat(int nid) {
    return stats[nid];
  }
  /*! \brief get leaf vector given nid */
  inline bst_float* leafvec(int nid) {
    if (leaf_vector.size() == 0) return NULL;
    return &leaf_vector[nid * param.size_leaf_vector];
  }
  /*! \brief get leaf vector given nid */
  inline const bst_float* leafvec(int nid) const{
    if (leaf_vector.size() == 0) return NULL;
    return &leaf_vector[nid * param.size_leaf_vector];
  }
  /*! \brief initialize the model */
  inline void InitModel(void) {
    param.num_nodes = param.num_roots;
    nodes.resize(param.num_nodes);
    stats.resize(param.num_nodes);
    leaf_vector.resize(param.num_nodes * param.size_leaf_vector, 0.0f);
    for (int i = 0; i < param.num_nodes; i ++) {
      nodes[i].set_leaf(0.0f);
      nodes[i].set_parent(-1);
    }
  }
  /*! 
   * \brief load model from stream
   * \param fi input stream
   */
  inline void LoadModel(utils::IStream &fi) {
    utils::Check(fi.Read(&param, sizeof(Param)) > 0,
                 "TreeModel: wrong format");
    nodes.resize(param.num_nodes); stats.resize(param.num_nodes);
    utils::Check(fi.Read(&nodes[0], sizeof(Node) * nodes.size()) > 0,
                 "TreeModel: wrong format");
    utils::Check(fi.Read(&stats[0], sizeof(NodeStat) * stats.size()) > 0,
                 "TreeModel: wrong format");
    if (param.size_leaf_vector != 0) {
      utils::Check(fi.Read(&leaf_vector), "TreeModel: wrong format");
    }
    // chg deleted nodes
    deleted_nodes.resize(0);
    for (int i = param.num_roots; i < param.num_nodes; i ++) {
      if (nodes[i].is_root()) deleted_nodes.push_back(i);
    }
    utils::Assert(static_cast<int>(deleted_nodes.size()) == param.num_deleted,
                  "number of deleted nodes do not match");
  }
  /*! 
   * \brief save model to stream
   * \param fo output stream
   */
  inline void SaveModel(utils::IStream &fo) const {
    utils::Assert(param.num_nodes == static_cast<int>(nodes.size()),
                  "Tree::SaveModel");
    utils::Assert(param.num_nodes == static_cast<int>(stats.size()),
                  "Tree::SaveModel");
    fo.Write(&param, sizeof(Param));
    fo.Write(&nodes[0], sizeof(Node) * nodes.size());
    fo.Write(&stats[0], sizeof(NodeStat) * nodes.size());
    if (param.size_leaf_vector != 0) fo.Write(leaf_vector);
  }
  /*! 
   * \brief add child nodes to node
   * \param nid node id to add childs
   */
  inline void AddChilds(int nid) {
    int pleft  = this->AllocNode();
    int pright = this->AllocNode();
    nodes[nid].cleft_  = pleft;
    nodes[nid].cright_ = pright;
    nodes[nodes[nid].cleft() ].set_parent(nid, true);
    nodes[nodes[nid].cright()].set_parent(nid, false);
  }
  /*! 
   * \brief only add a right child to a leaf node 
   * \param node id to add right child
   */
  inline void AddRightChild(int nid) {
    int pright = this->AllocNode();
    nodes[nid].right  = pright;
    nodes[nodes[nid].right].set_parent(nid, false);
  }
  /*!
   * \brief get current depth
   * \param nid node id
   * \param pass_rchild whether right child is not counted in depth
   */
  inline int GetDepth(int nid, bool pass_rchild = false) const {
    int depth = 0;
    while (!nodes[nid].is_root()) {
      if (!pass_rchild || nodes[nid].is_left_child()) ++depth;
      nid = nodes[nid].parent();
    }
    return depth;
  }
  /*!
   * \brief get maximum depth
   * \param nid node id
   */
  inline int MaxDepth(int nid) const {
    if (nodes[nid].is_leaf()) return 0;
    return std::max(MaxDepth(nodes[nid].cleft())+1,
                     MaxDepth(nodes[nid].cright())+1);
  }
  /*!
   * \brief get maximum depth
   */
  inline int MaxDepth(void) {
    int maxd = 0;
    for (int i = 0; i < param.num_roots; ++i) {
      maxd = std::max(maxd, MaxDepth(i));
    }
    return maxd;
  }
  /*! \brief number of extra nodes besides the root */
  inline int num_extra_nodes(void) const {
    return param.num_nodes - param.num_roots - param.num_deleted;
  }
  /*! 
   * \brief dump model to text string
   * \param fmap feature map of feature types
   * \param with_stats whether dump out statistics as well
   * \return the string of dumped model
   */
  inline std::string DumpModel(const utils::FeatMap& fmap, bool with_stats) {
    std::stringstream fo("");
    for (int i = 0; i < param.num_roots; ++i) {
      this->Dump(i, fo, fmap, 0, with_stats);
    }
    return fo.str();
  }

 private:
  void Dump(int nid, std::stringstream &fo,
            const utils::FeatMap& fmap, int depth, bool with_stats) {
    for (int i = 0;  i < depth; ++i) {
      fo << '\t';
    }
    if (nodes[nid].is_leaf()) {
      fo << nid << ":leaf=" << nodes[nid].leaf_value();
      if (with_stats) {
        stat(nid).Print(fo, true);
      }
      fo << '\n';
    } else {
      // right then left,
      TSplitCond cond = nodes[nid].split_cond();
      const unsigned split_index = nodes[nid].split_index();
      if (split_index < fmap.size()) {
        switch (fmap.type(split_index)) {
          case utils::FeatMap::kIndicator: {
            int nyes = nodes[nid].default_left() ?
                nodes[nid].cright() : nodes[nid].cleft();
            fo << nid << ":[" << fmap.name(split_index) << "] yes=" << nyes
               << ",no=" << nodes[nid].cdefault();
            break;
          }
          case utils::FeatMap::kInteger: {
            fo << nid << ":[" << fmap.name(split_index) << "<"
               << int(float(cond)+1.0f)
               << "] yes=" << nodes[nid].cleft()
               << ",no=" << nodes[nid].cright()
               << ",missing=" << nodes[nid].cdefault();
            break;
          }
          case utils::FeatMap::kFloat:
          case utils::FeatMap::kQuantitive: {
            fo << nid << ":[" << fmap.name(split_index) << "<"<< float(cond)
               << "] yes=" << nodes[nid].cleft()
               << ",no=" << nodes[nid].cright()
               << ",missing=" << nodes[nid].cdefault();
            break;
          }
          default: utils::Error("unknown fmap type");
        }
      } else {
        fo << nid << ":[f" << split_index << "<"<< float(cond)
           << "] yes=" << nodes[nid].cleft()
           << ",no=" << nodes[nid].cright()
           << ",missing=" << nodes[nid].cdefault();
      }
      if (with_stats) {
        fo << ' ';
        stat(nid).Print(fo, false);
      }
      fo << '\n';
      this->Dump(nodes[nid].cleft(), fo, fmap, depth+1, with_stats);
      this->Dump(nodes[nid].cright(), fo, fmap, depth+1, with_stats);
    }
  }
};

/*! \brief node statistics used in regression tree */
struct RTreeNodeStat {
  /*! \brief loss chg caused by current split */
  float loss_chg;
  /*! \brief sum of hessian values, used to measure coverage of data */
  float sum_hess;
  /*! \brief weight of current node */
  float base_weight;
  /*! \brief number of child that is leaf node known up to now */
  int   leaf_child_cnt;
  /*! \brief print information of current stats to fo */
  inline void Print(std::stringstream &fo, bool is_leaf) const {
    if (!is_leaf) {
      fo << "gain=" << loss_chg << ",cover=" << sum_hess;
    } else {
      fo << "cover=" << sum_hess;
    }
  }
};

/*! \brief define regression tree to be the most common tree model */
class RegTree: public TreeModel<bst_float, RTreeNodeStat>{
 public:
  /*! 
   * \brief dense feature vector that can be taken by RegTree
   * to do tranverse efficiently
   * and can be construct from sparse feature vector
   */
  struct FVec {
    /*! 
     * \brief a union value of value and flag
     * when flag == -1, this indicate the value is missing
     */
    union Entry{
      float fvalue;
      int flag;
    };
    std::vector<Entry> data;
    /*! \brief intialize the vector with size vector */
    inline void Init(size_t size) {
      Entry e; e.flag = -1;
      data.resize(size);
      std::fill(data.begin(), data.end(), e);
    }
    /*! \brief fill the vector with sparse vector */
    inline void Fill(const RowBatch::Inst &inst) {
      for (bst_uint i = 0; i < inst.length; ++i) {
        data[inst[i].index].fvalue = inst[i].fvalue;
      }
    }
    /*! \brief drop the trace after fill, must be called after fill */
    inline void Drop(const RowBatch::Inst &inst) {      
      for (bst_uint i = 0; i < inst.length; ++i) {
        data[inst[i].index].flag = -1;
      }
    }
    /*! \brief get ith value */
    inline float fvalue(size_t i) const {
      return data[i].fvalue;
    }
    /*! \brief check whether i-th entry is missing */
    inline bool is_missing(size_t i) const {
      return data[i].flag == -1;
    }
  };
  /*!
   * \brief get the leaf index 
   * \param feats dense feature vector, if the feature is missing the field is set to NaN
   * \param root_gid starting root index of the instance
   * \return the leaf index of the given feature 
   */
  inline int GetLeafIndex(const FVec&feat, unsigned root_id = 0) const {
    // start from groups that belongs to current data
    int pid = static_cast<int>(root_id);
    // tranverse tree
    while (!(*this)[ pid ].is_leaf()) {
      unsigned split_index = (*this)[pid].split_index();
      pid = this->GetNext(pid, feat.fvalue(split_index), feat.is_missing(split_index));
    }
    return pid;
  }
  /*!
   * \brief get the prediction of regression tree, only accepts dense feature vector
   * \param feats dense feature vector, if the feature is missing the field is set to NaN
   * \param root_gid starting root index of the instance
   * \return the leaf index of the given feature 
   */
  inline float Predict(const FVec &feat, unsigned root_id = 0) const {
    int pid = this->GetLeafIndex(feat, root_id);
    return (*this)[pid].leaf_value();
  }
  /*! \brief get next position of the tree given current pid */
  inline int GetNext(int pid, float fvalue, bool is_unknown) const {
    float split_value = (*this)[pid].split_cond();
    if (is_unknown) {
      return (*this)[pid].cdefault();
    } else {
      if (fvalue < split_value) {
        return (*this)[pid].cleft();
      } else {
        return (*this)[pid].cright();
      }
    }
  }
};

}  // namespace tree
}  // namespace xgboost
#endif  // XGBOOST_TREE_MODEL_H_