import warnings
from string import ascii_letters
from itertools import product
from functools import partial

import numpy as np
from pandas import (DataFrame, Series, MultiIndex, date_range, period_range,
                    TimeGrouper, Categorical)
import pandas.util.testing as tm

from .pandas_vb_common import setup  # noqa


method_blacklist = {
    'object': {'median', 'prod', 'sem', 'cumsum', 'sum', 'cummin', 'mean',
               'max', 'skew', 'cumprod', 'cummax', 'rank', 'pct_change', 'min',
               'var', 'mad', 'describe', 'std'},
    'datetime': {'median', 'prod', 'sem', 'cumsum', 'sum', 'mean', 'skew',
                 'cumprod', 'cummax', 'pct_change', 'var', 'mad', 'describe',
                 'std'}
}


class ApplyDictReturn(object):
    goal_time = 0.2

    def setup(self):
        self.labels = np.arange(1000).repeat(10)
        self.data = Series(np.random.randn(len(self.labels)))

    def time_groupby_apply_dict_return(self):
        self.data.groupby(self.labels).apply(lambda x: {'first': x.values[0],
                                                        'last': x.values[-1]})


class Apply(object):

    goal_time = 0.2

    def setup_cache(self):
        N = 10**4
        labels = np.random.randint(0, 2000, size=N)
        labels2 = np.random.randint(0, 3, size=N)
        df = DataFrame({'key': labels,
                        'key2': labels2,
                        'value1': np.random.randn(N),
                        'value2': ['foo', 'bar', 'baz', 'qux'] * (N // 4)
                        })
        return df

    def time_scalar_function_multi_col(self, df):
        df.groupby(['key', 'key2']).apply(lambda x: 1)

    def time_scalar_function_single_col(self, df):
        df.groupby('key').apply(lambda x: 1)

    @staticmethod
    def df_copy_function(g):
        # ensure that the group name is available (see GH #15062)
        g.name
        return g.copy()

    def time_copy_function_multi_col(self, df):
        df.groupby(['key', 'key2']).apply(self.df_copy_function)

    def time_copy_overhead_single_col(self, df):
        df.groupby('key').apply(self.df_copy_function)


class Groups(object):

    goal_time = 0.2

    param_names = ['key']
    params = ['int64_small', 'int64_large', 'object_small', 'object_large']

    def setup_cache(self):
        size = 10**6
        data = {'int64_small': Series(np.random.randint(0, 100, size=size)),
                'int64_large': Series(np.random.randint(0, 10000, size=size)),
                'object_small': Series(
                    tm.makeStringIndex(100).take(
                        np.random.randint(0, 100, size=size))),
                'object_large': Series(
                    tm.makeStringIndex(10000).take(
                        np.random.randint(0, 10000, size=size)))}
        return data

    def setup(self, data, key):
        self.ser = data[key]

    def time_series_groups(self, data, key):
        self.ser.groupby(self.ser).groups


class GroupManyLabels(object):

    goal_time = 0.2
    params = [1, 1000]
    param_names = ['ncols']

    def setup(self, ncols):
        N = 1000
        data = np.random.randn(N, ncols)
        self.labels = np.random.randint(0, 100, size=N)
        self.df = DataFrame(data)

    def time_sum(self, ncols):
        self.df.groupby(self.labels).sum()


class Nth(object):

    goal_time = 0.2

    param_names = ['dtype']
    params = ['float32', 'float64', 'datetime', 'object']

    def setup(self, dtype):
        N = 10**5
        # with datetimes (GH7555)
        if dtype == 'datetime':
            values = date_range('1/1/2011', periods=N, freq='s')
        elif dtype == 'object':
            values = ['foo'] * N
        else:
            values = np.arange(N).astype(dtype)

        key = np.arange(N)
        self.df = DataFrame({'key': key, 'values': values})
        self.df.iloc[1, 1] = np.nan  # insert missing data

    def time_frame_nth_any(self, dtype):
        self.df.groupby('key').nth(0, dropna='any')

    def time_groupby_nth_all(self, dtype):
        self.df.groupby('key').nth(0, dropna='all')

    def time_frame_nth(self, dtype):
        self.df.groupby('key').nth(0)

    def time_series_nth_any(self, dtype):
        self.df['values'].groupby(self.df['key']).nth(0, dropna='any')

    def time_groupby_nth_all(self, dtype):
        self.df['values'].groupby(self.df['key']).nth(0, dropna='all')

    def time_series_nth(self, dtype):
        self.df['values'].groupby(self.df['key']).nth(0)


class DateAttributes(object):

    goal_time = 0.2

    def setup(self):
        rng = date_range('1/1/2000', '12/31/2005', freq='H')
        self.year, self.month, self.day = rng.year, rng.month, rng.day
        self.ts = Series(np.random.randn(len(rng)), index=rng)

    def time_len_groupby_object(self):
        len(self.ts.groupby([self.year, self.month, self.day]))


class Int64(object):

    goal_time = 0.2

    def setup(self):
        arr = np.random.randint(-1 << 12, 1 << 12, (1 << 17, 5))
        i = np.random.choice(len(arr), len(arr) * 5)
        arr = np.vstack((arr, arr[i]))
        i = np.random.permutation(len(arr))
        arr = arr[i]
        self.cols = list('abcde')
        self.df = DataFrame(arr, columns=self.cols)
        self.df['jim'], self.df['joe'] = np.random.randn(2, len(self.df)) * 10

    def time_overflow(self):
        self.df.groupby(self.cols).max()


class CountMultiDtype(object):

    goal_time = 0.2

    def setup_cache(self):
        n = 10000
        offsets = np.random.randint(n, size=n).astype('timedelta64[ns]')
        dates = np.datetime64('now') + offsets
        dates[np.random.rand(n) > 0.5] = np.datetime64('nat')
        offsets[np.random.rand(n) > 0.5] = np.timedelta64('nat')
        value2 = np.random.randn(n)
        value2[np.random.rand(n) > 0.5] = np.nan
        obj = np.random.choice(list('ab'), size=n).astype(object)
        obj[np.random.randn(n) > 0.5] = np.nan
        df = DataFrame({'key1': np.random.randint(0, 500, size=n),
                        'key2': np.random.randint(0, 100, size=n),
                        'dates': dates,
                        'value2': value2,
                        'value3': np.random.randn(n),
                        'ints': np.random.randint(0, 1000, size=n),
                        'obj': obj,
                        'offsets': offsets})
        return df

    def time_multi_count(self, df):
        df.groupby(['key1', 'key2']).count()


class CountMultiInt(object):

    goal_time = 0.2

    def setup_cache(self):
        n = 10000
        df = DataFrame({'key1': np.random.randint(0, 500, size=n),
                        'key2': np.random.randint(0, 100, size=n),
                        'ints': np.random.randint(0, 1000, size=n),
                        'ints2': np.random.randint(0, 1000, size=n)})
        return df

    def time_multi_int_count(self, df):
        df.groupby(['key1', 'key2']).count()

    def time_multi_int_nunique(self, df):
        df.groupby(['key1', 'key2']).nunique()


class AggFunctions(object):

    goal_time = 0.2

    def setup_cache():
        N = 10**5
        fac1 = np.array(['A', 'B', 'C'], dtype='O')
        fac2 = np.array(['one', 'two'], dtype='O')
        df = DataFrame({'key1': fac1.take(np.random.randint(0, 3, size=N)),
                        'key2': fac2.take(np.random.randint(0, 2, size=N)),
                        'value1': np.random.randn(N),
                        'value2': np.random.randn(N),
                        'value3': np.random.randn(N)})
        return df

    def time_different_str_functions(self, df):
        df.groupby(['key1', 'key2']).agg({'value1': 'mean',
                                          'value2': 'var',
                                          'value3': 'sum'})

    def time_different_numpy_functions(self, df):
        df.groupby(['key1', 'key2']).agg({'value1': np.mean,
                                          'value2': np.var,
                                          'value3': np.sum})

    def time_different_python_functions_multicol(self, df):
        df.groupby(['key1', 'key2']).agg([sum, min, max])

    def time_different_python_functions_singlecol(self, df):
        df.groupby('key1').agg([sum, min, max])


class GroupStrings(object):

    goal_time = 0.2

    def setup(self):
        n = 2 * 10**5
        alpha = list(map(''.join, product(ascii_letters, repeat=4)))
        data = np.random.choice(alpha, (n // 5, 4), replace=False)
        data = np.repeat(data, 5, axis=0)
        self.df = DataFrame(data, columns=list('abcd'))
        self.df['joe'] = (np.random.randn(len(self.df)) * 10).round(3)
        self.df = self.df.sample(frac=1).reset_index(drop=True)

    def time_multi_columns(self):
        self.df.groupby(list('abcd')).max()


class MultiColumn(object):

    goal_time = 0.2

    def setup_cache(self):
        N = 10**5
        key1 = np.tile(np.arange(100, dtype=object), 1000)
        key2 = key1.copy()
        np.random.shuffle(key1)
        np.random.shuffle(key2)
        df = DataFrame({'key1': key1,
                        'key2': key2,
                        'data1': np.random.randn(N),
                        'data2': np.random.randn(N)})
        return df

    def time_lambda_sum(self, df):
        df.groupby(['key1', 'key2']).agg(lambda x: x.values.sum())

    def time_cython_sum(self, df):
        df.groupby(['key1', 'key2']).sum()

    def time_col_select_lambda_sum(self, df):
        df.groupby(['key1', 'key2'])['data1'].agg(lambda x: x.values.sum())

    def time_col_select_numpy_sum(self, df):
        df.groupby(['key1', 'key2'])['data1'].agg(np.sum)


class Size(object):

    goal_time = 0.2

    def setup(self):
        n = 10**5
        offsets = np.random.randint(n, size=n).astype('timedelta64[ns]')
        dates = np.datetime64('now') + offsets
        self.df = DataFrame({'key1': np.random.randint(0, 500, size=n),
                             'key2': np.random.randint(0, 100, size=n),
                             'value1': np.random.randn(n),
                             'value2': np.random.randn(n),
                             'value3': np.random.randn(n),
                             'dates': dates})
        self.draws = Series(np.random.randn(n))
        labels = Series(['foo', 'bar', 'baz', 'qux'] * (n // 4))
        self.cats = labels.astype('category')

    def time_multi_size(self):
        self.df.groupby(['key1', 'key2']).size()

    def time_dt_timegrouper_size(self):
        with warnings.catch_warnings(record=True):
            self.df.groupby(TimeGrouper(key='dates', freq='M')).size()

    def time_category_size(self):
        self.draws.groupby(self.cats).size()


class GroupByMethods(object):

    goal_time = 0.2

    param_names = ['dtype', 'method', 'application']
    params = [['int', 'float', 'object', 'datetime'],
              ['all', 'any', 'bfill', 'count', 'cumcount', 'cummax', 'cummin',
               'cumprod', 'cumsum', 'describe', 'ffill', 'first', 'head',
               'last', 'mad', 'max', 'min', 'median', 'mean', 'nunique',
               'pct_change', 'prod', 'rank', 'sem', 'shift', 'size', 'skew',
               'std', 'sum', 'tail', 'unique', 'value_counts', 'var'],
              ['direct', 'transformation']]

    def setup(self, dtype, method, application):
        if method in method_blacklist.get(dtype, {}):
            raise NotImplementedError  # skip benchmark
        ngroups = 1000
        size = ngroups * 2
        rng = np.arange(ngroups)
        values = rng.take(np.random.randint(0, ngroups, size=size))
        if dtype == 'int':
            key = np.random.randint(0, size, size=size)
        elif dtype == 'float':
            key = np.concatenate([np.random.random(ngroups) * 0.1,
                                  np.random.random(ngroups) * 10.0])
        elif dtype == 'object':
            key = ['foo'] * size
        elif dtype == 'datetime':
            key = date_range('1/1/2011', periods=size, freq='s')

        df = DataFrame({'values': values, 'key': key})

        if application == 'transform':
            if method == 'describe':
                raise NotImplementedError

            self.as_group_method = lambda: df.groupby(
                'key')['values'].transform(method)
            self.as_field_method = lambda: df.groupby(
                'values')['key'].transform(method)
        else:
            self.as_group_method = getattr(df.groupby('key')['values'], method)
            self.as_field_method = getattr(df.groupby('values')['key'], method)

    def time_dtype_as_group(self, dtype, method, application):
        self.as_group_method()

    def time_dtype_as_field(self, dtype, method, application):
        self.as_field_method()


class Float32(object):
    # GH 13335
    goal_time = 0.2

    def setup(self):
        tmp1 = (np.random.random(10000) * 0.1).astype(np.float32)
        tmp2 = (np.random.random(10000) * 10.0).astype(np.float32)
        tmp = np.concatenate((tmp1, tmp2))
        arr = np.repeat(tmp, 10)
        self.df = DataFrame(dict(a=arr, b=arr))

    def time_sum(self):
        self.df.groupby(['a'])['b'].sum()


class Categories(object):

    goal_time = 0.2

    def setup(self):
        N = 10**5
        arr = np.random.random(N)
        data = {'a': Categorical(np.random.randint(10000, size=N)),
                'b': arr}
        self.df = DataFrame(data)
        data = {'a': Categorical(np.random.randint(10000, size=N),
                                 ordered=True),
                'b': arr}
        self.df_ordered = DataFrame(data)
        data = {'a': Categorical(np.random.randint(100, size=N),
                                 categories=np.arange(10000)),
                'b': arr}
        self.df_extra_cat = DataFrame(data)

    def time_groupby_sort(self):
        self.df.groupby('a')['b'].count()

    def time_groupby_nosort(self):
        self.df.groupby('a', sort=False)['b'].count()

    def time_groupby_ordered_sort(self):
        self.df_ordered.groupby('a')['b'].count()

    def time_groupby_ordered_nosort(self):
        self.df_ordered.groupby('a', sort=False)['b'].count()

    def time_groupby_extra_cat_sort(self):
        self.df_extra_cat.groupby('a')['b'].count()

    def time_groupby_extra_cat_nosort(self):
        self.df_extra_cat.groupby('a', sort=False)['b'].count()


class Datelike(object):
    # GH 14338
    goal_time = 0.2
    params = ['period_range', 'date_range', 'date_range_tz']
    param_names = ['grouper']

    def setup(self, grouper):
        N = 10**4
        rng_map = {'period_range': period_range,
                   'date_range': date_range,
                   'date_range_tz': partial(date_range, tz='US/Central')}
        self.grouper = rng_map[grouper]('1900-01-01', freq='D', periods=N)
        self.df = DataFrame(np.random.randn(10**4, 2))

    def time_sum(self, grouper):
        self.df.groupby(self.grouper).sum()


class SumBools(object):
    # GH 2692
    goal_time = 0.2

    def setup(self):
        N = 500
        self.df = DataFrame({'ii': range(N),
                             'bb': [True] * N})

    def time_groupby_sum_booleans(self):
        self.df.groupby('ii').sum()


class SumMultiLevel(object):
    # GH 9049
    goal_time = 0.2
    timeout = 120.0

    def setup(self):
        N = 50
        self.df = DataFrame({'A': list(range(N)) * 2,
                             'B': range(N * 2),
                             'C': 1}).set_index(['A', 'B'])

    def time_groupby_sum_multiindex(self):
        self.df.groupby(level=[0, 1]).sum()


class Transform(object):

    goal_time = 0.2

    def setup(self):
        n1 = 400
        n2 = 250
        index = MultiIndex(levels=[np.arange(n1), tm.makeStringIndex(n2)],
                           labels=[np.repeat(range(n1), n2).tolist(),
                                   list(range(n2)) * n1],
                           names=['lev1', 'lev2'])
        arr = np.random.randn(n1 * n2, 3)
        arr[::10000, 0] = np.nan
        arr[1::10000, 1] = np.nan
        arr[2::10000, 2] = np.nan
        data = DataFrame(arr, index=index, columns=['col1', 'col20', 'col3'])
        self.df = data

        n = 20000
        self.df1 = DataFrame(np.random.randint(1, n, (n, 3)),
                             columns=['jim', 'joe', 'jolie'])
        self.df2 = self.df1.copy()
        self.df2['jim'] = self.df2['joe']

        self.df3 = DataFrame(np.random.randint(1, (n / 10), (n, 3)),
                             columns=['jim', 'joe', 'jolie'])
        self.df4 = self.df3.copy()
        self.df4['jim'] = self.df4['joe']

    def time_transform_lambda_max(self):
        self.df.groupby(level='lev1').transform(lambda x: max(x))

    def time_transform_ufunc_max(self):
        self.df.groupby(level='lev1').transform(np.max)

    def time_transform_multi_key1(self):
        self.df1.groupby(['jim', 'joe'])['jolie'].transform('max')

    def time_transform_multi_key2(self):
        self.df2.groupby(['jim', 'joe'])['jolie'].transform('max')

    def time_transform_multi_key3(self):
        self.df3.groupby(['jim', 'joe'])['jolie'].transform('max')

    def time_transform_multi_key4(self):
        self.df4.groupby(['jim', 'joe'])['jolie'].transform('max')


class TransformBools(object):

    goal_time = 0.2

    def setup(self):
        N = 120000
        transition_points = np.sort(np.random.choice(np.arange(N), 1400))
        transitions = np.zeros(N, dtype=np.bool)
        transitions[transition_points] = True
        self.g = transitions.cumsum()
        self.df = DataFrame({'signal': np.random.rand(N)})

    def time_transform_mean(self):
        self.df['signal'].groupby(self.g).transform(np.mean)


class TransformNaN(object):
    # GH 12737
    goal_time = 0.2

    def setup(self):
        self.df_nans = DataFrame({'key': np.repeat(np.arange(1000), 10),
                                  'B': np.nan,
                                  'C': np.nan})
        self.df_nans.loc[4::10, 'B':'C'] = 5

    def time_first(self):
        self.df_nans.groupby('key').transform('first')