# 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.abc import Iterable

from pandas.core import groupby as pd_groupby

from ... import opcodes
from ...core import OutputType
from ...serialize import AnyField
from ..operands import DataFrameOperandMixin, DataFrameOperand
from ..utils import parse_index


class GroupByIndex(DataFrameOperandMixin, DataFrameOperand):
    _op_type_ = opcodes.INDEX
    _op_module_ = 'dataframe.groupby'

    _selection = AnyField('selection')

    def __init__(self, selection=None, output_types=None, **kw):
        super().__init__(_selection=selection, _output_types=output_types, **kw)

    @property
    def selection(self):
        return self._selection

    @property
    def groupby_params(self):
        params = self.inputs[0].op.groupby_params
        params['selection'] = self.selection
        return params

    def build_mock_groupby(self, **kwargs):
        groupby_op = self.inputs[0].op
        return groupby_op.build_mock_groupby(**kwargs)[self.selection]

    def __call__(self, groupby):
        indexed = groupby.op.build_mock_groupby()[self.selection]

        if isinstance(indexed, pd_groupby.SeriesGroupBy):
            self.output_types = [OutputType.series_groupby]
            params = dict(shape=(groupby.shape[0],), name=self.selection,
                          dtype=groupby.dtypes[self.selection],
                          index_value=groupby.index_value,
                          key_dtypes=groupby.key_dtypes)
        else:
            self.output_types = [OutputType.dataframe_groupby]

            if isinstance(self.selection, Iterable) and not isinstance(self.selection, str):
                item_list = list(self.selection)
            else:
                item_list = [self.selection]

            params = groupby.params.copy()
            params['dtypes'] = new_dtypes = groupby.dtypes[item_list]
            params['selection'] = self.selection
            params['shape'] = (groupby.shape[0], len(item_list))
            params['columns_value'] = parse_index(new_dtypes.index, store_data=True)

        return self.new_tileable([groupby], **params)

    @classmethod
    def tile(cls, op):
        in_groupby = op.inputs[0]
        out_groupby = op.outputs[0]

        chunks = []
        for c in in_groupby.chunks:
            if op.output_types[0] == OutputType.series_groupby:
                params = dict(shape=(c.shape[0],), name=op.selection,
                              index=(c.index[0],), dtype=c.dtypes[op.selection],
                              index_value=c.index_value, key_dtypes=c.key_dtypes)
            else:
                params = c.params.copy()
                params['dtypes'] = out_groupby.dtypes
                params['selection'] = op.selection
                params['shape'] = (c.shape[0], len(op.selection))

            new_op = op.copy().reset_key()
            chunks.append(new_op.new_chunk([c], **params))

        new_op = op.copy().reset_key()
        params = out_groupby.params.copy()
        new_nsplits = (in_groupby.nsplits[0], (len(op.selection),)) if out_groupby.ndim == 2 \
            else (in_groupby.nsplits[0],)
        params.update(dict(chunks=chunks, nsplits=new_nsplits))
        return new_op.new_tileables([in_groupby], **params)

    @classmethod
    def execute(cls, ctx, op):
        in_data = ctx[op.inputs[0].key]
        ctx[op.outputs[0].key] = in_data[op.selection]


def df_groupby_getitem(df_groupby, item):
    try:
        hash(item)
        hashable = True
    except TypeError:
        hashable = False

    if hashable and item in df_groupby.dtypes:
        output_types = [OutputType.series_groupby]
    elif isinstance(item, Iterable) and all(it in df_groupby.dtypes for it in item):
        output_types = [OutputType.dataframe_groupby]
    else:
        raise NameError(f'Cannot slice groupby with {item!r}')

    if df_groupby.selection:
        raise IndexError(f'Column(s) {df_groupby.selection!r} already selected')

    op = GroupByIndex(selection=item, output_types=output_types)
    return op(df_groupby)
