# Path: bartpy/model.py
# class Model:
# 
#     def __init__(self,
#                  data: Optional[Data],
#                  sigma: Sigma,
#                  trees: Optional[List[Tree]]=None,
#                  n_trees: int=50,
#                  alpha: float=0.95,
#                  beta: float=2.,
#                  k: int=2.,
#                  initializer: Initializer=SklearnTreeInitializer()):
# 
#         self.data = deepcopy(data)
#         self.alpha = float(alpha)
#         self.beta = float(beta)
#         self.k = k
#         self._sigma = sigma
#         self._prediction = None
#         self._initializer = initializer
# 
#         if trees is None:
#             self.n_trees = n_trees
#             self._trees = self.initialize_trees()
#             if self._initializer is not None:
#                 self._initializer.initialize_trees(self.refreshed_trees())
#         else:
#             self.n_trees = len(trees)
#             self._trees = trees
# 
#     def initialize_trees(self) -> List[Tree]:
#         trees = [Tree([LeafNode(Split(deepcopy(self.data)))]) for _ in range(self.n_trees)]
#         for tree in trees:
#             tree.update_y(tree.update_y(self.data.y.values / self.n_trees))
#         return trees
# 
#     def residuals(self) -> np.ndarray:
#         return self.data.y.values - self.predict()
# 
#     def unnormalized_residuals(self) -> np.ndarray:
#         return self.data.y.unnormalized_y - self.data.y.unnormalize_y(self.predict())
# 
#     def predict(self, X: np.ndarray=None) -> np.ndarray:
#         if X is not None:
#             return self._out_of_sample_predict(X)
#         return np.sum([tree.predict() for tree in self.trees], axis=0)
# 
#     def _out_of_sample_predict(self, X: np.ndarray) -> np.ndarray:
#         if type(X) == pd.DataFrame:
#             X: pd.DataFrame = X
#             X = X.values
#         return np.sum([tree.predict(X) for tree in self.trees], axis=0)
# 
#     @property
#     def trees(self) -> List[Tree]:
#         return self._trees
# 
#     def refreshed_trees(self) -> Generator[Tree, None, None]:
#         if self._prediction is None:
#             self._prediction = self.predict()
#         for tree in self._trees:
#             self._prediction -= tree.predict()
#             tree.update_y(self.data.y.values - self._prediction)
#             yield tree
#             self._prediction += tree.predict()
# 
#     @property
#     def sigma_m(self) -> float:
#         return 0.5 / (self.k * np.power(self.n_trees, 0.5))
# 
#     @property
#     def sigma(self) -> Sigma:
#         return self._sigma
# 
# Path: bartpy/node.py
# class LeafNode(TreeNode):
#     """
#     A representation of a leaf node in the tree
#     In addition to the normal work of a `Node`, a `LeafNode` is responsible for:
#         - Interacting with `Data`
#         - Making predictions
#     """
# 
#     def __init__(self, split: Split, depth=0, value=0.0):
#         self._value = value
#         super().__init__(split, depth, None, None)
# 
#     def set_value(self, value: float) -> None:
#         self._value = value
# 
#     @property
#     def current_value(self):
#         return self._value
# 
#     def predict(self) -> float:
#         return self.current_value
# 
#     def is_splittable(self) -> bool:
#         return self.data.X.is_at_least_one_splittable_variable()
# 
# Path: bartpy/samplers/sampler.py
# class Sampler(ABC):
# 
#     @abstractmethod
#     def step(self, model: Model, tree: Tree) -> bool:
#         raise NotImplementedError()
# 
# Path: bartpy/samplers/scalar.py
# class NormalScalarSampler():
# 
#     def __init__(self,
#                  cache_size: int=1000):
#         self._cache_size = cache_size
#         self._cache = []
# 
#     def sample(self):
#         if len(self._cache) == 0:
#             self.refresh_cache()
#         return self._cache.pop()
# 
#     def refresh_cache(self):
#         self._cache = list(np.random.normal(size=self._cache_size))


import numpy as np
from bartpy.model import Model
from bartpy.node import LeafNode
from bartpy.samplers.sampler import Sampler
from bartpy.samplers.scalar import NormalScalarSampler



class LeafNodeSampler(Sampler):
    """
    Responsible for generating samples of the leaf node predictions
    Essentially just draws from a normal distribution with prior specified by model parameters

    Uses a cache of draws from a normal(0, 1) distribution to improve sampling performance
    """

    def __init__(self,
                 scalar_sampler=NormalScalarSampler(60000)):
        self._scalar_sampler = scalar_sampler


# Path: bartpy/model.py
# class Model:
# 
#     def __init__(self,
#                  data: Optional[Data],
#                  sigma: Sigma,
#                  trees: Optional[List[Tree]]=None,
#                  n_trees: int=50,
#                  alpha: float=0.95,
#                  beta: float=2.,
#                  k: int=2.,
#                  initializer: Initializer=SklearnTreeInitializer()):
# 
#         self.data = deepcopy(data)
#         self.alpha = float(alpha)
#         self.beta = float(beta)
#         self.k = k
#         self._sigma = sigma
#         self._prediction = None
#         self._initializer = initializer
# 
#         if trees is None:
#             self.n_trees = n_trees
#             self._trees = self.initialize_trees()
#             if self._initializer is not None:
#                 self._initializer.initialize_trees(self.refreshed_trees())
#         else:
#             self.n_trees = len(trees)
#             self._trees = trees
# 
#     def initialize_trees(self) -> List[Tree]:
#         trees = [Tree([LeafNode(Split(deepcopy(self.data)))]) for _ in range(self.n_trees)]
#         for tree in trees:
#             tree.update_y(tree.update_y(self.data.y.values / self.n_trees))
#         return trees
# 
#     def residuals(self) -> np.ndarray:
#         return self.data.y.values - self.predict()
# 
#     def unnormalized_residuals(self) -> np.ndarray:
#         return self.data.y.unnormalized_y - self.data.y.unnormalize_y(self.predict())
# 
#     def predict(self, X: np.ndarray=None) -> np.ndarray:
#         if X is not None:
#             return self._out_of_sample_predict(X)
#         return np.sum([tree.predict() for tree in self.trees], axis=0)
# 
#     def _out_of_sample_predict(self, X: np.ndarray) -> np.ndarray:
#         if type(X) == pd.DataFrame:
#             X: pd.DataFrame = X
#             X = X.values
#         return np.sum([tree.predict(X) for tree in self.trees], axis=0)
# 
#     @property
#     def trees(self) -> List[Tree]:
#         return self._trees
# 
#     def refreshed_trees(self) -> Generator[Tree, None, None]:
#         if self._prediction is None:
#             self._prediction = self.predict()
#         for tree in self._trees:
#             self._prediction -= tree.predict()
#             tree.update_y(self.data.y.values - self._prediction)
#             yield tree
#             self._prediction += tree.predict()
# 
#     @property
#     def sigma_m(self) -> float:
#         return 0.5 / (self.k * np.power(self.n_trees, 0.5))
# 
#     @property
#     def sigma(self) -> Sigma:
#         return self._sigma
# 
# Path: bartpy/node.py
# class LeafNode(TreeNode):
#     """
#     A representation of a leaf node in the tree
#     In addition to the normal work of a `Node`, a `LeafNode` is responsible for:
#         - Interacting with `Data`
#         - Making predictions
#     """
# 
#     def __init__(self, split: Split, depth=0, value=0.0):
#         self._value = value
#         super().__init__(split, depth, None, None)
# 
#     def set_value(self, value: float) -> None:
#         self._value = value
# 
#     @property
#     def current_value(self):
#         return self._value
# 
#     def predict(self) -> float:
#         return self.current_value
# 
#     def is_splittable(self) -> bool:
#         return self.data.X.is_at_least_one_splittable_variable()
# 
# Path: bartpy/samplers/sampler.py
# class Sampler(ABC):
# 
#     @abstractmethod
#     def step(self, model: Model, tree: Tree) -> bool:
#         raise NotImplementedError()
# 
# Path: bartpy/samplers/scalar.py
# class NormalScalarSampler():
# 
#     def __init__(self,
#                  cache_size: int=1000):
#         self._cache_size = cache_size
#         self._cache = []
# 
#     def sample(self):
#         if len(self._cache) == 0:
#             self.refresh_cache()
#         return self._cache.pop()
# 
#     def refresh_cache(self):
#         self._cache = list(np.random.normal(size=self._cache_size))


# Path: bartpy/samplers/leafnode.py
import numpy as np
from bartpy.model import Model
from bartpy.node import LeafNode
from bartpy.samplers.sampler import Sampler
from bartpy.samplers.scalar import NormalScalarSampler




class LeafNodeSampler(Sampler):
    """
    Responsible for generating samples of the leaf node predictions
    Essentially just draws from a normal distribution with prior specified by model parameters

    Uses a cache of draws from a normal(0, 1) distribution to improve sampling performance
    """

    def __init__(self,
                 scalar_sampler=NormalScalarSampler(60000)):
        self._scalar_sampler = scalar_sampler


    def step(self, model: Model, node: LeafNode) -> float: