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.
  
  from collections import OrderedDict
  from functools import reduce
  
  import pandas as pd
  
  from ..core import FuseChunkData, FuseChunk, Base, Entity, OutputType
  from ..operands import Operand, TileableOperandMixin, MapReduceOperand, Fuse
  from ..operands import ShuffleProxy, FuseChunkMixin
  from ..tensor.core import TENSOR_TYPE
  from ..tensor.operands import TensorOperandMixin
  from ..utils import calc_nsplits
  from .core import DATAFRAME_CHUNK_TYPE, SERIES_CHUNK_TYPE, DATAFRAME_TYPE, SERIES_TYPE, \
      INDEX_TYPE, DATAFRAME_GROUPBY_TYPE, SERIES_GROUPBY_TYPE, CATEGORICAL_TYPE
  from .utils import parse_index
  
  
  class DataFrameOperandMixin(TileableOperandMixin):
      __slots__ = ()
      _op_module_ = 'dataframe'
  
      def new_dataframes(self, inputs, shape=None, dtypes=None, index_value=None, columns_value=None,
                         chunks=None, nsplits=None, output_limit=None, kws=None, **kw):
          setattr(self, '_output_types', [OutputType.dataframe])
          return self.new_tileables(inputs, shape=shape, dtypes=dtypes, index_value=index_value,
                                    columns_value=columns_value, chunks=chunks, nsplits=nsplits,
                                    output_limit=output_limit, kws=kws, **kw)
  
      def new_dataframe(self, inputs, shape=None, dtypes=None, index_value=None, columns_value=None, **kw):
          if getattr(self, 'output_limit') != 1:
              raise TypeError('cannot new DataFrame with more than 1 outputs')
  
          return self.new_dataframes(inputs, shape=shape, dtypes=dtypes,
                                     index_value=index_value, columns_value=columns_value, **kw)[0]
  
      def new_seriess(self, inputs, shape=None, dtype=None, index_value=None, name=None,
                      chunks=None, nsplits=None, output_limit=None, kws=None, **kw):
          setattr(self, '_output_types', [OutputType.series])
          return self.new_tileables(inputs, shape=shape, dtype=dtype, index_value=index_value,
                                    name=name, chunks=chunks, nsplits=nsplits,
                                    output_limit=output_limit, kws=kws, **kw)
  
      def new_series(self, inputs, shape=None, dtype=None, index_value=None, name=None, **kw):
          if getattr(self, 'output_limit') != 1:
              raise TypeError('cannot new Series with more than 1 outputs')
  
          return self.new_seriess(inputs, shape=shape, dtype=dtype,
                                  index_value=index_value, name=name, **kw)[0]
  
      def new_indexes(self, inputs, shape=None, dtype=None, index_value=None, name=None,
                      chunks=None, nsplits=None, output_limit=None, kws=None, **kw):
          setattr(self, '_output_types', [OutputType.index])
          return self.new_tileables(inputs, shape=shape, dtype=dtype, index_value=index_value,
                                    name=name, chunks=chunks, nsplits=nsplits,
                                    output_limit=output_limit, kws=kws, **kw)
  
      def new_index(self, inputs, shape=None, dtype=None, index_value=None, name=None, **kw):
          if getattr(self, 'output_limit') != 1:
              raise TypeError('cannot new Index with more than 1 outputs')
  
          return self.new_indexes(inputs, shape=shape, dtype=dtype,
                                  index_value=index_value, name=name, **kw)[0]
  
      def new_scalars(self, inputs, dtype=None, chunks=None, output_limit=None, kws=None, **kw):
          setattr(self, '_output_types', [OutputType.scalar])
          return self.new_tileables(inputs, shape=(), dtype=dtype, chunks=chunks, nsplits=(),
                                    output_limit=output_limit, kws=kws, **kw)
  
      def new_scalar(self, inputs, dtype=None, **kw):
          if getattr(self, 'output_limit') != 1:
              raise TypeError('cannot new tensor with more than 1 outputs')
  
          return self.new_scalars(inputs, dtype=dtype, **kw)[0]
  
      def new_categoricals(self, inputs, shape=None, dtype=None, categories_value=None,
                           chunks=None, nsplits=None, output_limit=None, kws=None, **kw):
          setattr(self, '_output_types', [OutputType.categorical])
          return self.new_tileables(inputs, shape=shape, dtype=dtype,
                                    categories_value=categories_value, chunks=chunks,
                                    nsplits=nsplits, output_limit=output_limit,
                                    kws=kws, **kw)
  
      def new_categorical(self, inputs, shape=None, dtype=None, categories_value=None, **kw):
          if getattr(self, 'output_limit') != 1:
              raise TypeError('cannot new Categorical with more than 1 outputs')
  
          return self.new_categoricals(inputs, shape=shape, dtype=dtype,
                                       categories_value=categories_value, **kw)[0]
  
      @classmethod
      def _process_groupby_params(cls, groupby_params):
          new_groupby_params = groupby_params.copy()
          if isinstance(groupby_params['by'], list):
              by = []
              for v in groupby_params['by']:
                  if isinstance(v, (Base, Entity)):
                      by.append(cls.concat_tileable_chunks(v).chunks[0])
                  else:
                      by.append(v)
              new_groupby_params['by'] = by
          return new_groupby_params
  
      @classmethod
      def _get_groupby_inputs(cls, groupby, groupby_params):
          inputs = [groupby]
          chunk_inputs = list(groupby.chunks)
          if isinstance(groupby_params['by'], list):
              for chunk_v, v in zip(groupby_params['by'], groupby.op.groupby_params['by']):
                  if isinstance(v, (Base, Entity)):
                      inputs.append(v)
                      chunk_inputs.append(chunk_v)
          return inputs, chunk_inputs
  
      @classmethod
      def concat_tileable_chunks(cls, tileable):
          from .merge.concat import DataFrameConcat, GroupByConcat
  
          df = tileable
          assert not df.is_coarse()
  
          if isinstance(df, DATAFRAME_TYPE):
              chunk = DataFrameConcat(output_types=[OutputType.dataframe]).new_chunk(
                  df.chunks, shape=df.shape, index=(0, 0), dtypes=df.dtypes,
                  index_value=df.index_value, columns_value=df.columns_value)
              return DataFrameConcat(output_types=[OutputType.dataframe]).new_dataframe(
                  [df], shape=df.shape, chunks=[chunk],
                  nsplits=tuple((s,) for s in df.shape), dtypes=df.dtypes,
                  index_value=df.index_value, columns_value=df.columns_value)
          elif isinstance(df, SERIES_TYPE):
              chunk = DataFrameConcat(output_types=[OutputType.series]).new_chunk(
                  df.chunks, shape=df.shape, index=(0,), dtype=df.dtype,
                  index_value=df.index_value, name=df.name)
              return DataFrameConcat(output_types=[OutputType.series]).new_series(
                  [df], shape=df.shape, chunks=[chunk],
                  nsplits=tuple((s,) for s in df.shape), dtype=df.dtype,
                  index_value=df.index_value, name=df.name)
          elif isinstance(df, INDEX_TYPE):
              chunk = DataFrameConcat(output_types=[OutputType.index]).new_chunk(
                  df.chunks, shape=df.shape, index=(0,), dtype=df.dtype,
                  index_value=df.index_value, name=df.name)
              return DataFrameConcat(output_types=[OutputType.index]).new_series(
                  [df], shape=df.shape, chunks=[chunk],
                  nsplits=tuple((s,) for s in df.shape), dtype=df.dtype,
                  index_value=df.index_value, name=df.name)
          elif isinstance(df, (DATAFRAME_GROUPBY_TYPE, SERIES_GROUPBY_TYPE)):
              output_type = OutputType.dataframe_groupby \
                  if isinstance(df, DATAFRAME_GROUPBY_TYPE) else OutputType.series_groupby
              groupby_params = cls._process_groupby_params(df.op.groupby_params)
              inputs, chunk_inputs = cls._get_groupby_inputs(df, groupby_params)
              chunk = GroupByConcat(groups=df.chunks, groupby_params=groupby_params,
                                    output_types=[output_type]).new_chunk(
                  chunk_inputs, **df.params)
              return GroupByConcat(groups=[df], groupby_params=df.op.groupby_params,
                                   output_types=[output_type]).new_dataframe(
                  inputs, chunks=[chunk], **df.params)
          elif isinstance(df, CATEGORICAL_TYPE):
              chunk = DataFrameConcat(output_types=[OutputType.categorical]).new_chunk(
                  df.chunks, shape=df.shape, index=(0,), dtype=df.dtype,
                  categories_value=df.categories_value)
              return DataFrameConcat(output_types=[OutputType.categorical]).new_categorical(
                  [df], shape=df.shape, chunks=[chunk],
                  nsplits=tuple((s,) for s in df.shape), dtype=df.dtype,
                  categories_value=df.categories_value)
          elif isinstance(df, TENSOR_TYPE):
              return TensorOperandMixin.concat_tileable_chunks(tileable)
          else:
              raise NotImplementedError
  
      @classmethod
      def create_tileable_from_chunks(cls, chunks, inputs=None, **kw):
          ndim = chunks[0].ndim
          index_min, index_max = [None] * ndim, [None] * ndim
          for c in chunks:
              for ax, i in enumerate(c.index):
                  if index_min[ax] is None:
                      index_min[ax] = i
                  else:
                      index_min[ax] = min(i, index_min[ax])
                  if index_max[ax] is None:
                      index_max[ax] = i
                  else:
                      index_max[ax] = max(i, index_max[ax])
  
          # gen {chunk index -> shape}
          chunk_index_to_shape = OrderedDict()
          chunk_index_to_chunk = dict()
          for c in chunks:
              new_index = []
              for ax, i in enumerate(c.index):
                  new_index.append(i - index_min[ax])
              chunk_index_to_shape[tuple(new_index)] = c.shape
              chunk_index_to_chunk[tuple(new_index)] = c
  
          nsplits = calc_nsplits(chunk_index_to_shape)
          shape = tuple(sum(ns) for ns in nsplits)
          chunk_shape = tuple(len(ns) for ns in nsplits)
          op = chunks[0].op.copy().reset_key()
          if isinstance(chunks[0], DATAFRAME_CHUNK_TYPE):
              params = cls._calc_dataframe_params(chunk_index_to_chunk, chunk_shape)
              params.update(kw)
              return op.new_dataframe(inputs, shape=shape, chunks=chunks,
                                      nsplits=nsplits, **params)
          else:
              assert isinstance(chunks[0], SERIES_CHUNK_TYPE)
              params = cls._calc_series_params(chunks)
              params.update(kw)
              return op.new_series(inputs, shape=shape, chunks=chunks,
                                   nsplits=nsplits, **params)
  
      @classmethod
      def _calc_dataframe_params(cls, chunk_index_to_chunks, chunk_shape):
          dtypes = pd.concat([chunk_index_to_chunks[0, i].dtypes
                              for i in range(chunk_shape[1])])
          columns_value = parse_index(dtypes.index, store_data=True)
          pd_indxes = [chunk_index_to_chunks[i, 0].index_value.to_pandas()
                       for i in range(chunk_shape[0])]
          pd_index = reduce(lambda x, y: x.append(y), pd_indxes)
          index_value = parse_index(pd_index)
          return {'dtypes': dtypes, 'columns_value': columns_value,
                  'index_value': index_value}
  
      @classmethod
      def _calc_series_params(cls, chunks):
          pd_indexes = [c.index_value.to_pandas() for c in chunks]
          pd_index = reduce(lambda x, y: x.append(y), pd_indexes)
          index_value = parse_index(pd_index)
          return {'dtype': chunks[0].dtype, 'index_value': index_value}
  
      def get_fetch_op_cls(self, _):
          from ..operands import ShuffleProxy
          from .fetch import DataFrameFetchShuffle, DataFrameFetch
          if isinstance(self, ShuffleProxy):
              cls = DataFrameFetchShuffle
          else:
              cls = DataFrameFetch
  
          def _inner(**kw):
              return cls(output_types=self.output_types, **kw)
  
          return _inner
  
      def get_fuse_op_cls(self, _):
          return DataFrameFuseChunk
  
  
  DataFrameOperand = Operand
  
  
  class DataFrameShuffleProxy(ShuffleProxy, DataFrameOperandMixin):
      def __init__(self, sparse=None, output_types=None, **kwargs):
          super().__init__(_sparse=sparse, _output_types=output_types, **kwargs)
  
  
  DataFrameMapReduceOperand = MapReduceOperand
  
  
  class DataFrameFuseChunkMixin(FuseChunkMixin, DataFrameOperandMixin):
      __slots__ = ()
  
      def _create_chunk(self, output_idx, index, **kw):
          data = FuseChunkData(_index=index, _shape=kw.pop('shape', None), _op=self, **kw)
  
          return FuseChunk(data)
  
  
  class DataFrameFuseChunk(Fuse, DataFrameFuseChunkMixin):
      def __init__(self, sparse=False, **kwargs):
          super().__init__(_sparse=sparse, **kwargs)
  
      @property
      def output_types(self):
          return self._operands[-1].output_types
  