# Path: bartpy/tree.py
# class Tree:
#     """
#     An encapsulation of the structure of a single decision tree
#     Contains no logic, but keeps track of 4 different kinds of nodes within the tree:
#       - leaf nodes
#       - decision nodes
#       - splittable leaf nodes
#       - prunable decision nodes
# 
#     Parameters
#     ----------
#     nodes: List[Node]
#         All nodes contained in the tree, i.e. decision and leaf nodes
#     """
# 
#     def __init__(self, nodes: List[TreeNode]):
#         self._nodes = nodes
#         self.cache_up_to_date = False
#         self._prediction = None
# 
#     @property
#     def nodes(self) -> List[TreeNode]:
#         """
#         List of all nodes contained in the tree
#         """
#         return self._nodes
# 
#     @property
#     def leaf_nodes(self) -> List[LeafNode]:
#         """
#         List of all of the leaf nodes in the tree
#         """
#         return [x for x in self._nodes if type(x) == LeafNode]
# 
#     @property
#     def splittable_leaf_nodes(self) -> List[LeafNode]:
#         """
#         List of all leaf nodes in the tree which can be split in a non-degenerate way
#         i.e. not all rows of the covariate matrix are duplicates
#         """
#         return [x for x in self.leaf_nodes if x.is_splittable()]
# 
#     @property
#     def decision_nodes(self) -> List[DecisionNode]:
#         """
#         List of decision nodes in the tree.
#         Decision nodes are internal split nodes, i.e. not leaf nodes
#         """
#         return [x for x in self._nodes if type(x) == DecisionNode]
# 
#     @property
#     def prunable_decision_nodes(self) -> List[DecisionNode]:
#         """
#         List of decision nodes in the tree that are suitable for pruning
#         In particular, decision nodes that have two leaf node children
#         """
#         return [x for x in self.decision_nodes if x.is_prunable()]
# 
#     def update_y(self, y: np.ndarray) -> None:
#         """
#         Update the cached value of the target array in all nodes
#         Used to pass in the residuals from the sum of all of the other trees
#         """
#         self.cache_up_to_date = False
#         for node in self.nodes:
#             node.update_y(y)
# 
#     def predict(self, X: np.ndarray=None) -> np.ndarray:
#         """
#         Generate a set of predictions with the same dimensionality as the target array
#         Note that the prediction is from one tree, so represents only (1 / number_of_trees) of the target
#         """
#         if X is not None:
#             return self._out_of_sample_predict(X)
# 
#         if self.cache_up_to_date:
#             return self._prediction
#         for leaf in self.leaf_nodes:
#             if self._prediction is None:
#                 self._prediction = np.zeros(self.nodes[0].data.X.n_obsv)
#             self._prediction[leaf.split.condition()] = leaf.predict()
#         self.cache_up_to_date = True
#         return self._prediction
# 
#     def _out_of_sample_predict(self, X) -> np.ndarray:
#         """
#         Prediction for a covariate matrix not used for training
# 
#         Note that this is quite slow
# 
#         Parameters
#         ----------
#         X: pd.DataFrame
#             Covariates to predict for
#         Returns
#         -------
#         np.ndarray
#         """
#         prediction = np.array([0.] * len(X))
#         for leaf in self.leaf_nodes:
#             prediction[leaf.split.condition(X)] = leaf.predict()
#         return prediction
# 
#     def remove_node(self, node: TreeNode) -> None:
#         """
#         Remove a single node from the tree
#         Note that this is non-recursive, only drops the node and not any children
#         """
#         self._nodes.remove(node)
# 
#     def add_node(self, node: TreeNode) -> None:
#         """
#         Add a node to the tree
#         Note that this is non-recursive, only adds the node and not any children
#         """
#         self._nodes.append(node)


from typing import Generator
from bartpy.tree import Tree



class Initializer(object):
    """
    The abstract interface for the tree initializers.

    Initializers are responsible for setting the starting values of the model, in particular:
      - structure of decision and leaf nodes
      - variables and values used in splits
      - values of leaf nodes

    Good initialization of trees helps speed up convergence of sampling

    Default behaviour is to leave trees uninitialized
    """


# Path: bartpy/tree.py
# class Tree:
#     """
#     An encapsulation of the structure of a single decision tree
#     Contains no logic, but keeps track of 4 different kinds of nodes within the tree:
#       - leaf nodes
#       - decision nodes
#       - splittable leaf nodes
#       - prunable decision nodes
# 
#     Parameters
#     ----------
#     nodes: List[Node]
#         All nodes contained in the tree, i.e. decision and leaf nodes
#     """
# 
#     def __init__(self, nodes: List[TreeNode]):
#         self._nodes = nodes
#         self.cache_up_to_date = False
#         self._prediction = None
# 
#     @property
#     def nodes(self) -> List[TreeNode]:
#         """
#         List of all nodes contained in the tree
#         """
#         return self._nodes
# 
#     @property
#     def leaf_nodes(self) -> List[LeafNode]:
#         """
#         List of all of the leaf nodes in the tree
#         """
#         return [x for x in self._nodes if type(x) == LeafNode]
# 
#     @property
#     def splittable_leaf_nodes(self) -> List[LeafNode]:
#         """
#         List of all leaf nodes in the tree which can be split in a non-degenerate way
#         i.e. not all rows of the covariate matrix are duplicates
#         """
#         return [x for x in self.leaf_nodes if x.is_splittable()]
# 
#     @property
#     def decision_nodes(self) -> List[DecisionNode]:
#         """
#         List of decision nodes in the tree.
#         Decision nodes are internal split nodes, i.e. not leaf nodes
#         """
#         return [x for x in self._nodes if type(x) == DecisionNode]
# 
#     @property
#     def prunable_decision_nodes(self) -> List[DecisionNode]:
#         """
#         List of decision nodes in the tree that are suitable for pruning
#         In particular, decision nodes that have two leaf node children
#         """
#         return [x for x in self.decision_nodes if x.is_prunable()]
# 
#     def update_y(self, y: np.ndarray) -> None:
#         """
#         Update the cached value of the target array in all nodes
#         Used to pass in the residuals from the sum of all of the other trees
#         """
#         self.cache_up_to_date = False
#         for node in self.nodes:
#             node.update_y(y)
# 
#     def predict(self, X: np.ndarray=None) -> np.ndarray:
#         """
#         Generate a set of predictions with the same dimensionality as the target array
#         Note that the prediction is from one tree, so represents only (1 / number_of_trees) of the target
#         """
#         if X is not None:
#             return self._out_of_sample_predict(X)
# 
#         if self.cache_up_to_date:
#             return self._prediction
#         for leaf in self.leaf_nodes:
#             if self._prediction is None:
#                 self._prediction = np.zeros(self.nodes[0].data.X.n_obsv)
#             self._prediction[leaf.split.condition()] = leaf.predict()
#         self.cache_up_to_date = True
#         return self._prediction
# 
#     def _out_of_sample_predict(self, X) -> np.ndarray:
#         """
#         Prediction for a covariate matrix not used for training
# 
#         Note that this is quite slow
# 
#         Parameters
#         ----------
#         X: pd.DataFrame
#             Covariates to predict for
#         Returns
#         -------
#         np.ndarray
#         """
#         prediction = np.array([0.] * len(X))
#         for leaf in self.leaf_nodes:
#             prediction[leaf.split.condition(X)] = leaf.predict()
#         return prediction
# 
#     def remove_node(self, node: TreeNode) -> None:
#         """
#         Remove a single node from the tree
#         Note that this is non-recursive, only drops the node and not any children
#         """
#         self._nodes.remove(node)
# 
#     def add_node(self, node: TreeNode) -> None:
#         """
#         Add a node to the tree
#         Note that this is non-recursive, only adds the node and not any children
#         """
#         self._nodes.append(node)


# Path: bartpy/initializers/initializer.py
from typing import Generator
from bartpy.tree import Tree




class Initializer(object):
    """
    The abstract interface for the tree initializers.

    Initializers are responsible for setting the starting values of the model, in particular:
      - structure of decision and leaf nodes
      - variables and values used in splits
      - values of leaf nodes

    Good initialization of trees helps speed up convergence of sampling

    Default behaviour is to leave trees uninitialized
    """


    def initialize_tree(self, tree: Tree) -> None: