# Path: bartpy/data.py
# def make_bartpy_data(X: Union[np.ndarray, pd.DataFrame],
#                      y: np.ndarray,
#                      normalize: bool=True) -> 'Data':
#     X = format_covariate_matrix(X)
#     y = y.astype(float)
#     return Data(X, y, normalize=normalize)
# 
# Path: bartpy/samplers/unconstrainedtree/proposer.py
# def uniformly_sample_grow_mutation(tree: Tree) -> TreeMutation:
#     node = random_splittable_leaf_node(tree)
#     updated_node = sample_split_node(node)
#     return GrowMutation(node, updated_node)
# 
# def uniformly_sample_prune_mutation(tree: Tree) -> TreeMutation:
#     node = random_prunable_decision_node(tree)
#     updated_node = LeafNode(node.split, depth=node.depth)
#     return PruneMutation(node, updated_node)
# 
# Path: bartpy/split.py
# class Split:
#     """
#     The Split class represents the conditioned data at any point in the decision tree
#     It contains the logic for:
# 
#      - Maintaining a record of which rows of the covariate matrix are in the split
#      - Being able to easily access a `Data` object with the relevant rows
#      - Applying `SplitConditions` to further break up the data
#     """
# 
#     def __init__(self, data: Data, combined_condition: Optional[CombinedCondition]=None):
#         self._data = data
#         if combined_condition is None:
#             combined_condition = CombinedCondition(self._data.X.variables, [])
#         self._combined_condition = combined_condition
# 
#     @property
#     def data(self):
#         return self._data
# 
#     def combined_condition(self):
#         return self._combined_condition
# 
#     def condition(self, X: np.ndarray=None) -> np.array:
#         if X is None:
#             return ~self._data.mask
#         else:
#             return self.out_of_sample_condition(X)
# 
#     def out_of_sample_condition(self, X: np.ndarray) -> np.ndarray:
#         return self._combined_condition.condition(X)
# 
#     def out_of_sample_conditioner(self) -> CombinedCondition:
#         return self._combined_condition
# 
#     def __add__(self, other: SplitCondition):
# 
#         return Split(self._data + other,
#                      self._combined_condition + other)
# 
#     def most_recent_split_condition(self) -> Optional[SplitCondition]:
#         return self._combined_condition.most_recent_split_condition()
# 
# Path: bartpy/tree.py
# class Tree:
#     def __init__(self, nodes: List[TreeNode]):
#     def nodes(self) -> List[TreeNode]:
#     def leaf_nodes(self) -> List[LeafNode]:
#     def splittable_leaf_nodes(self) -> List[LeafNode]:
#     def decision_nodes(self) -> List[DecisionNode]:
#     def prunable_decision_nodes(self) -> List[DecisionNode]:
#     def update_y(self, y: np.ndarray) -> None:
#     def predict(self, X: np.ndarray=None) -> np.ndarray:
#     def _out_of_sample_predict(self, X) -> np.ndarray:
#     def remove_node(self, node: TreeNode) -> None:
#     def add_node(self, node: TreeNode) -> None:
# def mutate(tree: Tree, mutation: TreeMutation) -> None:
# def deep_copy_tree(tree: Tree):


import unittest
import pandas as pd
import numpy as np
from bartpy.data import make_bartpy_data
from bartpy.samplers.unconstrainedtree.proposer import uniformly_sample_grow_mutation, uniformly_sample_prune_mutation
from bartpy.split import Split
from bartpy.tree import LeafNode, Tree, DecisionNode




class TestPruneTreeMutationProposer(unittest.TestCase):

    def setUp(self):
        self.data = make_bartpy_data(pd.DataFrame({"a": [1, 2]}), np.array([1, 2]), normalize=False)

# Path: bartpy/data.py
# def make_bartpy_data(X: Union[np.ndarray, pd.DataFrame],
#                      y: np.ndarray,
#                      normalize: bool=True) -> 'Data':
#     X = format_covariate_matrix(X)
#     y = y.astype(float)
#     return Data(X, y, normalize=normalize)
# 
# Path: bartpy/samplers/unconstrainedtree/proposer.py
# def uniformly_sample_grow_mutation(tree: Tree) -> TreeMutation:
#     node = random_splittable_leaf_node(tree)
#     updated_node = sample_split_node(node)
#     return GrowMutation(node, updated_node)
# 
# def uniformly_sample_prune_mutation(tree: Tree) -> TreeMutation:
#     node = random_prunable_decision_node(tree)
#     updated_node = LeafNode(node.split, depth=node.depth)
#     return PruneMutation(node, updated_node)
# 
# Path: bartpy/split.py
# class Split:
#     """
#     The Split class represents the conditioned data at any point in the decision tree
#     It contains the logic for:
# 
#      - Maintaining a record of which rows of the covariate matrix are in the split
#      - Being able to easily access a `Data` object with the relevant rows
#      - Applying `SplitConditions` to further break up the data
#     """
# 
#     def __init__(self, data: Data, combined_condition: Optional[CombinedCondition]=None):
#         self._data = data
#         if combined_condition is None:
#             combined_condition = CombinedCondition(self._data.X.variables, [])
#         self._combined_condition = combined_condition
# 
#     @property
#     def data(self):
#         return self._data
# 
#     def combined_condition(self):
#         return self._combined_condition
# 
#     def condition(self, X: np.ndarray=None) -> np.array:
#         if X is None:
#             return ~self._data.mask
#         else:
#             return self.out_of_sample_condition(X)
# 
#     def out_of_sample_condition(self, X: np.ndarray) -> np.ndarray:
#         return self._combined_condition.condition(X)
# 
#     def out_of_sample_conditioner(self) -> CombinedCondition:
#         return self._combined_condition
# 
#     def __add__(self, other: SplitCondition):
# 
#         return Split(self._data + other,
#                      self._combined_condition + other)
# 
#     def most_recent_split_condition(self) -> Optional[SplitCondition]:
#         return self._combined_condition.most_recent_split_condition()
# 
# Path: bartpy/tree.py
# class Tree:
#     def __init__(self, nodes: List[TreeNode]):
#     def nodes(self) -> List[TreeNode]:
#     def leaf_nodes(self) -> List[LeafNode]:
#     def splittable_leaf_nodes(self) -> List[LeafNode]:
#     def decision_nodes(self) -> List[DecisionNode]:
#     def prunable_decision_nodes(self) -> List[DecisionNode]:
#     def update_y(self, y: np.ndarray) -> None:
#     def predict(self, X: np.ndarray=None) -> np.ndarray:
#     def _out_of_sample_predict(self, X) -> np.ndarray:
#     def remove_node(self, node: TreeNode) -> None:
#     def add_node(self, node: TreeNode) -> None:
# def mutate(tree: Tree, mutation: TreeMutation) -> None:
# def deep_copy_tree(tree: Tree):


# Path: tests/test_proposer.py
import unittest
import pandas as pd
import numpy as np
from bartpy.data import make_bartpy_data
from bartpy.samplers.unconstrainedtree.proposer import uniformly_sample_grow_mutation, uniformly_sample_prune_mutation
from bartpy.split import Split
from bartpy.tree import LeafNode, Tree, DecisionNode





class TestPruneTreeMutationProposer(unittest.TestCase):

    def setUp(self):
        self.data = make_bartpy_data(pd.DataFrame({"a": [1, 2]}), np.array([1, 2]), normalize=False)

        self.d = LeafNode(Split(self.data))