# Copyright 1999-2021 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 sys
from collections.abc import Iterable

import cloudpickle
import numpy as np
import pandas as pd
from pandas.core.groupby import DataFrameGroupBy, SeriesGroupBy

from .version import parse as parse_version

_HAS_SQUEEZE = parse_version(pd.__version__) < parse_version('1.1.0')


class GroupByWrapper:
    def __init__(self, obj, groupby_obj=None, keys=None, axis=0, level=None, grouper=None,
                 exclusions=None, selection=None, as_index=True, sort=True,
                 group_keys=True, squeeze=False, observed=False, mutated=False,
                 grouper_cache=None):

        def fill_value(v, key):
            return v if v is not None or groupby_obj is None else getattr(groupby_obj, key)

        self.obj = obj
        self.keys = fill_value(keys, 'keys')
        self.axis = fill_value(axis, 'axis')
        self.level = fill_value(level, 'level')
        self.exclusions = fill_value(exclusions, 'exclusions')
        self.selection = selection
        self.as_index = fill_value(as_index, 'as_index')
        self.sort = fill_value(sort, 'sort')
        self.group_keys = fill_value(group_keys, 'group_keys')
        self.squeeze = fill_value(squeeze, 'squeeze')
        self.observed = fill_value(observed, 'observed')
        self.mutated = fill_value(mutated, 'mutated')

        if groupby_obj is None:
            groupby_kw = dict(keys=keys, axis=axis, level=level, grouper=grouper,
                              exclusions=exclusions, as_index=as_index, group_keys=group_keys,
                              squeeze=squeeze, observed=observed, mutated=mutated)
            if not _HAS_SQUEEZE:  # pragma: no branch
                groupby_kw.pop('squeeze')

            if obj.ndim == 2:
                self.groupby_obj = DataFrameGroupBy(obj, **groupby_kw)
            else:
                self.groupby_obj = SeriesGroupBy(obj, **groupby_kw)
        else:
            self.groupby_obj = groupby_obj

        if grouper_cache:
            self.groupby_obj.grouper._cache = grouper_cache
        if selection:
            self.groupby_obj = self.groupby_obj[selection]

        self.is_frame = isinstance(self.groupby_obj, DataFrameGroupBy)

    def __getitem__(self, item):
        return GroupByWrapper(
            self.obj, keys=self.keys, axis=self.axis, level=self.level,
            grouper=self.groupby_obj.grouper, exclusions=self.exclusions, selection=item,
            as_index=self.as_index, sort=self.sort, group_keys=self.group_keys,
            squeeze=self.squeeze, observed=self.observed, mutated=self.mutated)

    def __getattr__(self, item):
        if item.startswith('_'):  # pragma: no cover
            return object.__getattribute__(self, item)
        if item in getattr(self.obj, 'columns', ()):
            return self.__getitem__(item)
        return getattr(self.groupby_obj, item)

    def __iter__(self):
        return self.groupby_obj.__iter__()

    def __sizeof__(self):
        return sys.getsizeof(self.obj) \
            + sys.getsizeof(getattr(self.groupby_obj.grouper, '_cache', None))

    def __reduce__(self):
        return type(self).from_tuple, (self.to_tuple(pickle_function=True, truncate=True),)

    def __bool__(self):
        return bool(np.prod(self.shape))

    @property
    def empty(self):
        return self.obj.empty

    @property
    def shape(self):
        shape = list(self.groupby_obj.obj.shape)
        if self.is_frame and self.selection:
            shape[1] = len(self.selection)
        return tuple(shape)

    @property
    def _selected_obj(self):
        return getattr(self.groupby_obj, '_selected_obj')

    def to_tuple(self, truncate=False, pickle_function=False):
        if self.selection and truncate:
            if isinstance(self.selection, Iterable) and not isinstance(self.selection, str):
                item_list = list(self.selection)
            else:
                item_list = [self.selection]
            item_set = set(item_list)

            if isinstance(self.keys, list):
                sel_keys = self.keys
            elif self.keys in self.obj.columns:
                sel_keys = [self.keys]
            else:
                sel_keys = []

            all_items = item_list + [k for k in sel_keys or () if k not in item_set]
            if set(all_items) == set(self.obj.columns):
                obj = self.obj
            else:
                obj = self.obj[all_items]
        else:
            obj = self.obj

        if pickle_function and callable(self.keys):
            keys = cloudpickle.dumps(self.keys)
        else:
            keys = self.keys

        return obj, keys, self.axis, self.level, self.exclusions, self.selection, \
            self.as_index, self.sort, self.group_keys, self.squeeze, self.observed, \
            self.mutated, getattr(self.groupby_obj.grouper, '_cache', dict())

    @classmethod
    def from_tuple(cls, tp):
        obj, keys, axis, level, exclusions, selection, as_index, sort, group_keys, squeeze, \
            observed, mutated, grouper_cache = tp

        if isinstance(keys, (bytes, bytearray)):
            keys = cloudpickle.loads(keys)

        return cls(obj, keys=keys, axis=axis, level=level, exclusions=exclusions, selection=selection,
                   as_index=as_index, sort=sort, group_keys=group_keys, squeeze=squeeze, observed=observed,
                   mutated=mutated, grouper_cache=grouper_cache)


def wrapped_groupby(obj, by=None, axis=0, level=None, as_index=True, sort=True, group_keys=True,
                    squeeze=False, observed=False):
    groupby_kw = dict(by=by, axis=axis, level=level, as_index=as_index, sort=sort,
                      group_keys=group_keys, squeeze=squeeze, observed=observed)
    if not _HAS_SQUEEZE:  # pragma: no branch
        groupby_kw.pop('squeeze')

    groupby_obj = obj.groupby(**groupby_kw)
    return GroupByWrapper(obj, groupby_obj=groupby_obj)
