# Path: bartpy/data.py
# class Data(object):
#     """
#     Encapsulates the data within a split of feature space.
#     Primarily used to cache computations on the data for better performance
# 
#     Parameters
#     ----------
#     X: np.ndarray
#         The subset of the covariate matrix that falls into the split
#     y: np.ndarray
#         The subset of the target array that falls into the split
#     normalize: bool
#         Whether to map the target into -0.5, 0.5
#     cache: bool
#         Whether to cache common values.
#         You really only want to turn this off if you're not going to the resulting object for anything (e.g. when testing)
#     """
# 
#     def __init__(self,
#                  X: np.ndarray,
#                  y: np.ndarray,
#                  mask: Optional[np.ndarray]=None,
#                  normalize: bool=False,
#                  unique_columns: List[int]=None,
#                  splittable_variables: Optional[List[Optional[bool]]]=None,
#                  y_sum: float=None,
#                  n_obsv: int=None):
# 
#         if mask is None:
#             mask = np.zeros_like(y).astype(bool)
#         self._mask: np.ndarray = mask
# 
#         if n_obsv is None:
#             n_obsv = (~self.mask).astype(int).sum()
# 
#         self._n_obsv = n_obsv
# 
#         self._X = CovariateMatrix(X, mask, n_obsv, unique_columns, splittable_variables)
#         self._y = Target(y, mask, n_obsv, normalize, y_sum)
# 
#     @property
#     def y(self) -> Target:
#         return self._y
# 
#     @property
#     def X(self) -> CovariateMatrix:
#         return self._X
# 
#     @property
#     def mask(self) -> np.ndarray:
#         return self._mask
# 
#     def update_y(self, y: np.ndarray) -> None:
#         self._y.update_y(y)
# 
#     def __add__(self, other: SplitCondition) -> 'Data':
#         updated_mask = self.X.update_mask(other)
# 
#         return Data(self.X.values,
#                     self.y.values,
#                     updated_mask,
#                     normalize=False,
#                     unique_columns=self._X._unique_columns,
#                     splittable_variables=self._X._splittable_variables,
#                     y_sum=other.carry_y_sum,
#                     n_obsv=other.carry_n_obsv)
# 
# 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/split.py
# class Split:
#     def __init__(self, data: Data, combined_condition: Optional[CombinedCondition]=None):
#     def data(self):
#     def combined_condition(self):
#     def condition(self, X: np.ndarray=None) -> np.array:
#     def out_of_sample_condition(self, X: np.ndarray) -> np.ndarray:
#     def out_of_sample_conditioner(self) -> CombinedCondition:
#     def __add__(self, other: SplitCondition):
#     def most_recent_split_condition(self) -> Optional[SplitCondition]:


from operator import le, gt
from bartpy.data import Data, make_bartpy_data
from bartpy.split import SplitCondition, Split, CombinedCondition
import unittest
import pandas as pd
import numpy as np
    def test_combined_condition_data(self):
        data = make_bartpy_data(pd.DataFrame({"a": [1, 2, 3, 4]}).values, np.array([1, 2, 1, 1]))

        first_left_condition, first_right_condition = SplitCondition(0, 3, le), SplitCondition(0, 3, gt)
        second_left_condition, second_right_condition = SplitCondition(0, 1, le), SplitCondition(0, 1, gt)

        split = Split(data)
        updated_split = split + first_left_condition + second_right_condition
        conditioned_data = updated_split.data
        self.assertListEqual([2, 3], list(conditioned_data.X.get_column(0)))

    def test_most_recent_split(self):
        data = make_bartpy_data(pd.DataFrame({"a": [1, 2, 3, 4]}).values, np.array([1, 2, 1, 1]))

        first_left_condition, first_right_condition = SplitCondition(0, 3, le), SplitCondition(0, 3, gt)
        second_left_condition, second_right_condition = SplitCondition(0, 1, le), SplitCondition(0, 1, gt)

        split = Split(data)
        updated_split = split + first_left_condition + second_right_condition
        self.assertEqual((split + first_left_condition).most_recent_split_condition(), first_left_condition)
        self.assertEqual(updated_split.most_recent_split_condition(), second_right_condition)


class TestCombinedCondition(unittest.TestCase):

    def setUp(self):
        self.X = np.array([1, 2, 4, 6, 3, 5]).reshape(6, 1)

    def test_single_condition(self):
        condition = SplitCondition(0, 3, gt)

# Path: bartpy/data.py
# class Data(object):
#     """
#     Encapsulates the data within a split of feature space.
#     Primarily used to cache computations on the data for better performance
# 
#     Parameters
#     ----------
#     X: np.ndarray
#         The subset of the covariate matrix that falls into the split
#     y: np.ndarray
#         The subset of the target array that falls into the split
#     normalize: bool
#         Whether to map the target into -0.5, 0.5
#     cache: bool
#         Whether to cache common values.
#         You really only want to turn this off if you're not going to the resulting object for anything (e.g. when testing)
#     """
# 
#     def __init__(self,
#                  X: np.ndarray,
#                  y: np.ndarray,
#                  mask: Optional[np.ndarray]=None,
#                  normalize: bool=False,
#                  unique_columns: List[int]=None,
#                  splittable_variables: Optional[List[Optional[bool]]]=None,
#                  y_sum: float=None,
#                  n_obsv: int=None):
# 
#         if mask is None:
#             mask = np.zeros_like(y).astype(bool)
#         self._mask: np.ndarray = mask
# 
#         if n_obsv is None:
#             n_obsv = (~self.mask).astype(int).sum()
# 
#         self._n_obsv = n_obsv
# 
#         self._X = CovariateMatrix(X, mask, n_obsv, unique_columns, splittable_variables)
#         self._y = Target(y, mask, n_obsv, normalize, y_sum)
# 
#     @property
#     def y(self) -> Target:
#         return self._y
# 
#     @property
#     def X(self) -> CovariateMatrix:
#         return self._X
# 
#     @property
#     def mask(self) -> np.ndarray:
#         return self._mask
# 
#     def update_y(self, y: np.ndarray) -> None:
#         self._y.update_y(y)
# 
#     def __add__(self, other: SplitCondition) -> 'Data':
#         updated_mask = self.X.update_mask(other)
# 
#         return Data(self.X.values,
#                     self.y.values,
#                     updated_mask,
#                     normalize=False,
#                     unique_columns=self._X._unique_columns,
#                     splittable_variables=self._X._splittable_variables,
#                     y_sum=other.carry_y_sum,
#                     n_obsv=other.carry_n_obsv)
# 
# 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/split.py
# class Split:
#     def __init__(self, data: Data, combined_condition: Optional[CombinedCondition]=None):
#     def data(self):
#     def combined_condition(self):
#     def condition(self, X: np.ndarray=None) -> np.array:
#     def out_of_sample_condition(self, X: np.ndarray) -> np.ndarray:
#     def out_of_sample_conditioner(self) -> CombinedCondition:
#     def __add__(self, other: SplitCondition):
#     def most_recent_split_condition(self) -> Optional[SplitCondition]:


# Path: tests/test_split.py
from operator import le, gt
from bartpy.data import Data, make_bartpy_data
from bartpy.split import SplitCondition, Split, CombinedCondition
import unittest
import pandas as pd
import numpy as np

    def test_combined_condition_data(self):
        data = make_bartpy_data(pd.DataFrame({"a": [1, 2, 3, 4]}).values, np.array([1, 2, 1, 1]))

        first_left_condition, first_right_condition = SplitCondition(0, 3, le), SplitCondition(0, 3, gt)
        second_left_condition, second_right_condition = SplitCondition(0, 1, le), SplitCondition(0, 1, gt)

        split = Split(data)
        updated_split = split + first_left_condition + second_right_condition
        conditioned_data = updated_split.data
        self.assertListEqual([2, 3], list(conditioned_data.X.get_column(0)))

    def test_most_recent_split(self):
        data = make_bartpy_data(pd.DataFrame({"a": [1, 2, 3, 4]}).values, np.array([1, 2, 1, 1]))

        first_left_condition, first_right_condition = SplitCondition(0, 3, le), SplitCondition(0, 3, gt)
        second_left_condition, second_right_condition = SplitCondition(0, 1, le), SplitCondition(0, 1, gt)

        split = Split(data)
        updated_split = split + first_left_condition + second_right_condition
        self.assertEqual((split + first_left_condition).most_recent_split_condition(), first_left_condition)
        self.assertEqual(updated_split.most_recent_split_condition(), second_right_condition)


class TestCombinedCondition(unittest.TestCase):

    def setUp(self):
        self.X = np.array([1, 2, 4, 6, 3, 5]).reshape(6, 1)

    def test_single_condition(self):
        condition = SplitCondition(0, 3, gt)

        combined_condition = CombinedCondition([0], [condition])