text: # Copyright 1999-2020 Alibaba Group Holding Ltd.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import cloudpickle
  import numpy as np
  
  from ...core import Object, OBJECT_TYPE, OBJECT_CHUNK_TYPE, recursive_tile
  from ...serialization.serializables import KeyField, Int32Field, DictField, AnyField, BoolField
  from ...tensor.core import TensorOrder
  from ...utils import has_unknown_shape, tokenize
  from ..operands import LearnOperand, LearnOperandMixin, OutputType
  
  
  class TreeObject(Object):
      def fetch(self, session=None, **kw):
          result = self._data.fetch(session=session, **kw)
          return cloudpickle.loads(result) \
              if isinstance(result, bytes) else result
  
  
  class TreeBase(LearnOperand, LearnOperandMixin):
      _input = KeyField('input')
      _leaf_size = Int32Field('leaf_size')
      _metric = AnyField('metric')
  
      _metric_params = DictField('metric_params')
  
      def __init__(self, leaf_size=None, metric=None,
                   metric_params=None, output_types=None, **kw):
          super().__init__(_leaf_size=leaf_size, _metric=metric,
                           _metric_params=metric_params,
                           _output_types=output_types, **kw)
          if self.output_types is None:
              self.output_types = [OutputType.object]
  
      @property
      def input(self):
          return self._input
  
      @property
      def leaf_size(self):
          return self._leaf_size
  
      @property
      def metric(self):
          return self._metric
  
      @property
      def metric_params(self):
          return self._metric_params
  
      def _set_inputs(self, inputs):
          super()._set_inputs(inputs)
          self._input = self._inputs[0]
  
      def __call__(self, a):
          return self.new_tileable([a])
  
      @classmethod
      def tile(cls, op):
          if has_unknown_shape(*op.inputs):
              yield
  
          # ball tree and kd tree requires the full data,
          # thus rechunk input tensor into 1 chunk
          inp = op.input.rechunk({ax: s for ax, s in enumerate(op.input.shape)})
          inp = yield from recursive_tile(inp)
          out = op.outputs[0]
  
          chunk_op = op.copy().reset_key()
          kw = out.params
          kw['index'] = inp.chunks[0].index
          chunk = chunk_op.new_chunk([inp.chunks[0]], kws=[kw])
  
          new_op = op.copy()
          tileable_kw = out.params
          tileable_kw['nsplits'] = ((1,),)
          tileable_kw['chunks'] = [chunk]
          return new_op.new_tileables(op.inputs, kws=[tileable_kw])
  
      @classmethod
      def execute(cls, ctx, op):
          if op.gpu:  # pragma: no cover
              raise NotImplementedError('Does not support tree-based '
                                        'nearest neighbors on GPU')
  
          a = ctx[op.input.key]
          tree = cls._tree_type(
              a, op.leaf_size, metric=op.metric,
              **(op.metric_params or dict()))
          ctx[op.outputs[0].key] = tree
  
  
  def _on_serialize_tree(tree):
      return cloudpickle.dumps(tree) if not hasattr(tree, 'key') else tree
  
  
  def _on_deserialize_tree(ser):
      return cloudpickle.loads(ser) if isinstance(ser, bytes) else ser
  
  
  class TreeQueryBase(LearnOperand, LearnOperandMixin):
      _input = KeyField('input')
      _tree = AnyField('tree', on_serialize=_on_serialize_tree,
                       on_deserialize=_on_deserialize_tree)
      _n_neighbors = Int32Field('n_neighbors')
      _return_distance = BoolField('return_distance')
  
      def __init__(self, tree=None, n_neighbors=None, return_distance=None,
                   output_types=None, **kw):
          super().__init__(_tree=tree, _n_neighbors=n_neighbors,
                           _return_distance=return_distance,
                           _output_types=output_types, **kw)
          if self.output_types is None:
              self.output_types = [OutputType.tensor] * self.output_limit
  
      @property
      def input(self):
          return self._input
  
      @property
      def tree(self):
          return self._tree
  
      @property
      def n_neighbors(self):
          return self._n_neighbors
  
      @property
      def return_distance(self):
          return self._return_distance
  
      @property
      def output_limit(self):
          return 2 if self._return_distance else 1
  
      def _set_inputs(self, inputs):
          super()._set_inputs(inputs)
          self._input = self._inputs[0]
          if isinstance(self._tree, (OBJECT_TYPE, OBJECT_CHUNK_TYPE)):
              self._tree = self._inputs[1]
  
      def _update_key(self):
          values = []
          for value in self._values_:
              if isinstance(value, self._tree_type):
                  values.append(cloudpickle.dumps(value))
              else:
                  values.append(value)
          self._obj_set('_key', tokenize(type(self).__name__, *values))
          return self
  
      def __call__(self, x):
          kws = []
          if self._return_distance:
              kws.append({'shape': (x.shape[0], self._n_neighbors),
                          'dtype': np.dtype(np.float64),
                          'order': x.order,
                          'type': 'distance'})
          kws.append({
              'shape': (x.shape[0], self._n_neighbors),
              'dtype': np.dtype(np.int64),
              'order': TensorOrder.C_ORDER,
              'type': 'indices'
          })
          inputs = [x]
          if isinstance(self._tree, OBJECT_TYPE):
              inputs.append(self._tree)
          return self.new_tileables(inputs, kws=kws, output_limit=len(kws))
  
      @classmethod
      def tile(cls, op):
          inp = op.input
  
          if inp.chunk_shape[1] != 1:
              if has_unknown_shape(inp):
                  yield
              inp = yield from recursive_tile(inp.rechunk({1: inp.shape[1]}))
  
          tree_chunk = None
          if isinstance(op.tree, OBJECT_TYPE):
              tree_chunk = op.tree.chunks[0]
          out_chunks = [[] for _ in range(len(op.outputs))]
          for chunk in inp.chunks:
              chunk_op = op.copy().reset_key()
              if tree_chunk is not None:
                  chunk_op._tree = tree_chunk
              chunk_kws = []
              if op.return_distance:
                  chunk_kws.append({
                      'shape': (chunk.shape[0], op.n_neighbors),
                      'dtype': np.dtype(np.float64),
                      'order': chunk.order,
                      'index': chunk.index,
                      'type': 'distance'
                  })
              chunk_kws.append({
                  'shape': (chunk.shape[0], op.n_neighbors),
                  'dtype': np.dtype(np.int64),
                  'order': TensorOrder.C_ORDER,
                  'index': chunk.index,
                  'type': 'indices'
              })
              chunk_inputs = [chunk]
              if tree_chunk is not None:
                  chunk_inputs.append(tree_chunk)
              chunks = chunk_op.new_chunks(chunk_inputs, kws=chunk_kws,
                                           output_limit=len(chunk_kws))
              for cs, c in zip(out_chunks, chunks):
                  cs.append(c)
  
          kws = [o.params for o in op.outputs]
          nsplits = list(inp.nsplits)
          nsplits[1] = (op.n_neighbors,)
          if op.return_distance:
              kws[0]['chunks'] = out_chunks[0]
              kws[0]['nsplits'] = tuple(nsplits)
          kws[-1]['chunks'] = out_chunks[-1]
          kws[-1]['nsplits'] = tuple(nsplits)
          new_op = op.copy()
          return new_op.new_tileables(op.inputs, kws=kws, output_limit=len(kws))
  
      @classmethod
      def execute(cls, ctx, op):
          if op.gpu:  # pragma: no cover
              raise NotImplementedError('Does not support tree-based '
                                        'nearest neighbors on GPU')
  
          x = ctx[op.input.key]
          if len(op.inputs) == 2:
              tree = ctx[op.tree.key]
          else:
              tree = op.tree
          tree = cloudpickle.loads(tree) if isinstance(tree, bytes) else tree
          ret = tree.query(x, op.n_neighbors, op.return_distance)
          if op.return_distance:
              ctx[op.outputs[0].key] = ret[0]
              ctx[op.outputs[1].key] = ret[1]
          else:
              ctx[op.outputs[0].key] = ret
  