text: import numpy as np
  import math
  import types
  
  from adiscr.node import *
  from adiscr.tree_vis import *
  
  
  def actions_per_level(level, dims):
      return 2**(level*dims)
  
  
  def compute_level(n, dims):
      actions = [1]
      level = 0
  
      while actions[level] < n:
          level += 1
          actions.append(actions[level-1]+actions_per_level(level, dims))
  
      if actions[level] == n:
          return level
      else:
          return level-1
  
  
  class Tree:
  
      def __init__(self, dims, size, error_function='direct'):
          """
          dims: integer > 0
          size: integer > 0
          error_function: ['direct', 'sqrt', 'square', 'cubic'] or a lambda function
          """
          assert dims > 0, "dims is < 1"
          assert size > 0, "size is < 1"
          assert (error_function in ['direct', 'sqrt', 'square', 'cubic']) or isinstance(
              error_function, types.FunctionType), "error_function has to be one of ['direct', 'sqrt', 'square', 'cubic'] or a function"
  
          self._dimensions = int(dims)
          self._size = int(size)
          Node._init_branch_matrix(self._dimensions)
  
          if error_function is 'direct':
              self._error_function = lambda d: d
          elif error_function is 'sqrt':
              self._error_function = lambda d: d**0.5
          elif error_function is 'square':
              self._error_function = lambda d: d**2
          elif error_function is 'cubic':
              self._error_function = lambda d: d**3
          else:
              self._error_function = error_function
  
          self._branch_factor = self._dimensions * 2
          root = Node(None, None, self._error_function, dims)
          self._nodes = [root]
          self._root = root
  
          init_level = compute_level(size, self._dimensions)
          for i in range(init_level):
              self.add_layer()
  
          self._total_distance = 0
          self._total_distance_count = 0
  
      def search_nearest_node(self, point, increase=True):
  
          point = self.correct_point(point.flatten())
  
          node, dist = self._root.search(point, increase)
  
          self._total_distance += dist
  
          self._total_distance_count += 1
  
          return node
  
      def update(self):
          _points_before = np.array(self.get_points())
  
          to_prune = self._prune_prospectives()
          pruned = 0
          for node in to_prune:
              node.delete()
              pruned += 1
  
          excess = self.get_current_size() - self.get_size()
          expanded = self._expand_usefull_nodes(pruned - excess)
  
          self._refresh_nodes()
          self._reset_values()
  
          _points_after = np.array(self.get_points())
  
          if pruned == expanded:
              for i in range(len(_points_after)):
                  if np.linalg.norm(_points_before[i] - _points_after[i]) > 0:
                      return True
              return False
          else:
              return True
  
      def feed(self, samples):
          for sample in samples:
              self.search_nearest_node(sample)
  
      def feed_and_update(self, samples):
          self.feed(samples)
          return self.update()
  
      def adapt_to_samples(self, samples, max_iterations=10):
          print("Adaption begun,", len(samples), "samples, max iterations", max_iterations)
          count = 0
          flag = True
          while flag and count < max_iterations:
              flag = self.feed_and_update(samples)
              print("Iteration", count, ", adapted:", not flag)
              count += 1
          return not flag
  
      def _prune_prospectives(self):
  
          nodes = self.get_prunable_nodes()
  
          mean_value = self.get_mean_value()
  
          result = []
  
          for node in nodes:
              estimated_future_value = node.get_value() + node.get_value_increase_if_cut()
              if(estimated_future_value < mean_value):
                  result.append(node)
  
          return result
  
      def _expand_usefull_nodes(self, n):
          nodes = sorted(self.get_nodes(recalculate=True), key=lambda node: node.get_value())
          suggestions = list(node.suggest_for_expand() for node in nodes)
  
          count_expantions = 0
          i = len(nodes)
          while count_expantions < n and i >= 0:
              i -= 1
              if(nodes[i].get_value() == 0):
                  continue
  
              to_expand = suggestions[i]
              new_nodes = []
              for suggestion in to_expand:
                  max_new_nodes = n - count_expantions - len(new_nodes)
                  new_nodes.extend(suggestion.expand(
                      nodes[i].get_location(), new_nodes_limit=max_new_nodes))
  
              self._nodes.extend(new_nodes)
              count_expantions += len(new_nodes)
  
          return count_expantions
  
      def get_node(self, index):
          node = self.get_nodes()[index]
          return node
  
      def recursive_traversal(self, func=(lambda node: node),
                              traverse_cond_func=(lambda node: True),
                              collect_cond_func=(lambda node: True)):
          res = []
          self._root.recursive_collection(res, func, traverse_cond_func, collect_cond_func)
          return res
  
      def get_values(self):
          return self.recursive_traversal(func=(lambda node: node.get_value()))
  
      def get_total_value(self):
          return np.sum(self.get_values())
  
      def get_prunable_nodes(self):
          return self.recursive_traversal(collect_cond_func=(lambda node: node.is_leaf()))
  
      def get_expendable_nodes(self):
          return self.recursive_traversal(collect_cond_func=(lambda node: node.is_expandable()))
  
      def get_mean_error(self):
          if self._total_distance_count == 0:
              return 0
  
          return np.sum(self.get_values()) / self._total_distance_count
  
      def get_mean_value(self):
          return np.sum(self.get_values()) / self.get_current_size()
  
      def _get_max_mean_distance(self):
          return 1 / (4 * self.get_current_size())
  
      def _reset_values(self):
          self.recursive_traversal(func=lambda node: node.reset_value())
          self._total_distance = 0
          self._total_distance_count = 0
  
      def add_layer(self):
          to_expand = self.get_expendable_nodes()
          for node in to_expand:
              new_nodes = node.expand()
              self._nodes.extend(new_nodes)
  
      def get_nodes(self, recalculate=False):
          if recalculate:
              self._nodes = self.recursive_traversal()
          return self._nodes
  
      def _refresh_nodes(self):
          self._nodes = self.recursive_traversal()
  
      def get_points(self):
          return np.array(list(node.get_location() for node in self.get_nodes()))
  
      def get_current_size(self):
          return len(self._nodes)
  
      def get_size(self):
          return self._size
  
      def print_all_nodes(self):
          nodes = self._nodes
          print('tree contains', len(nodes), 'nodes, min level=', self._min_level)
          for node in nodes:
              print(node)
  
      def plot(self, red_levels=False, save=False, filename='', plot_density_flag=False):
          plot(self, save=save, red_levels=red_levels,
               path=filename, plot_density_flag=plot_density_flag)
  
      @staticmethod
      def correct_point(point):
          new_point = []
          for c in point:
              if c > 1:
                  new_point.append(1)
              elif c < 0:
                  new_point.append(0)
              else:
                  new_point.append(c)
  
          return new_point
  