import numpy as np
import pandas as pd

np.random.seed(10)

df = pd.DataFrame({
    'key1': ['a', 'a', 'b', 'b', 'a'],
    'key2': ['one', 'two', 'one', 'two', 'one'],
    'data1': np.random.randn(5),
    'data2': np.random.randn(5)
})

df

# 1)
# - We want to compute the mean of the values in "data1" grouping by values of "key1"
# - groupby() computes the mapping between keys of the groups and rows of the dataframe
grouped = df['data1'].groupby(df['key1'])

# We are grouping a Series since we have a single column.
grouped

grouped2 = df[['data1', 'data2']].groupby(df['key1'])

# We are grouping a DataFrame since we have two columns.
grouped2

# For each value of the group we compute the mean of the corresponding rows.
grouped.mean()

# If we group by 2 keys, we end up with a hierarchical index Series.
means = df['data1'].groupby([df['key1'], df['key2']]).mean()

means

means.unstack()

df['data1']

# 2)
# We can also use arrays to infer the groups, as long as the size is
# the same as the number of rows.
states = np.array(['Ohio', 'California', 'California', 'Ohio', 'Ohio'])
years = np.array([2005, 2005, 2006, 2005, 2006])

print(len(states), len(years), len(df['data1']))

df['data1'].groupby([states, years]).mean()

# 3)
# The grouping information can be stored in the same data frame as the data
# Note that df['key2'] is excluded from the mean since it is not numerical.
df.groupby('key1').mean()

# Group by 2 keys the entire df.
df.groupby(['key1', 'key2']).mean()

# We can count the number of elements with size().
# Note that nan are excluded.
df.groupby(['key1', 'key2']).size()

df

for name, group in df.groupby('key1'):
    # group is the dataframe.
    print("\n# key=", name)
    print("group=\n", group)

# In case of grouping by multiple keys, the "key" is a tuple of values.
for (k1, k2), group in df.groupby(['key1', 'key2']):
    print("\n# key=", (k1, k2))
    print("group=\n", group)

# One can compute a dict out of the groupby in one line.
pieces = dict(list(df.groupby('key1')))

import pprint

pprint.pprint(pieces)

df

# Group by "key1".
grouped = df.groupby('key1')

# .groups.keys() to get the keys.
print("keys=", list(grouped.groups.keys()))

# A groupby object can be split by column after being computed.
print(df.groupby('key1')["data1"])

print(df.groupby('key1')["data1"].mean())
print(df.groupby('key1')["data2"].mean())

# It is equivalent to
# - "split and then select" and
# - "select and then split"
print(df.groupby('key1')["data1"].mean())

# Once we have selected "data1" there is no "key1" anymore so we use the array
# df["key1"] to label the values and group.
print(df["data1"].groupby(df['key1']).mean())

np.random.seed(10)

people = pd.DataFrame(
    np.random.randn(5, 5),
    columns=list("abcde"),
    index="Joe Steve Wes Jim Travis".split())

# Add NAs at row = 2 and columns = [1, 3]
people.iloc[2:3, [1, 3]] = np.nan

people

# Build a map from columns to group and aggregate.
mapping = {
    'a': 'red',
    'b': 'red',
    'c': 'blue',
    'd': 'blue',
    'e': 'red',
    'f': 'orange'
}

by_column = people.groupby(mapping, axis=1)

by_column.sum()

# Transform the dict into a fixed mapping series.
map_series = pd.Series(mapping)

print(map_series)

display(people.groupby(map_series, axis=1).sum())

# Group by length of name
people.groupby(len).sum()

columns = pd.MultiIndex.from_arrays(
    [['US', 'US', 'US', 'JP', 'JP'], [1, 3, 5, 1, 3]],
    names=['cty', 'tenor'])

columns

hier_df = pd.DataFrame(np.random.randn(4, 5), columns=columns)

hier_df

hier_df.groupby(level='cty', axis=1).count()

df

# split by values of key1 and compute quantile.
df.groupby('key1').quantile(0.9)

# One can use any custom function.
def peak_to_peak(arr):
    return arr.max() - arr.min()

df.groupby('key1').agg(peak_to_peak)

# Also functions like describe() work, although they are not
# aggregations.

df.groupby('key1')["data1"].describe()

df.groupby('key1').describe()

tips = pd.read_csv('~/src/github/pydata-book/examples/tips.csv')

tips['tip_pct'] = tips['tip'] / tips['total_bill']

tips.head()

grouped = tips.groupby(['day', 'smoker'])

# Select a column.
grouped_pct = grouped['tip_pct']
print("keys=", list(grouped_pct.groups.keys()))

# Equivalent.
grouped_pct.agg('mean')
grouped_pct.mean()

# Pass a list of aggregation functions.
funcs = ['mean', 'std', peak_to_peak]
grouped_pct.agg(funcs)

# Assign name to the functions.
funcs = [
    ('foo', 'mean'),
    ('bar', np.std)
]
grouped_pct.agg(funcs)

funcs = ['count', 'mean', 'max']
result = grouped['tip_pct', 'total_bill'].agg(funcs)

# It has hierarchical columns for both rows and columns.
result

# Aggregation functions can also be specified by dict.
funcs = {'tip': np.max, 'size': 'sum'}
grouped.agg(funcs)

funcs = {'tip_pct': ['min', 'max', 'mean', 'std'], 'size': 'sum'}
grouped.agg(funcs)

tips.groupby(['day', 'smoker'], as_index=True).mean()

# Returning a hierarchical index can be disabled.
tips.groupby(['day', 'smoker'], as_index=False).mean()

# This is equivalent to call reset_index().

tips.groupby(['day', 'smoker'], as_index=True).mean().reset_index()

# You want to select the top five tip_pct values by group.

def top(df, n=2, column='tip_pct'):
    return df.sort_values(by=column)[-n:]

top(tips, n=6)

# top() is called on each row group and then results
# are concat with pandas.concat, using labels from group name.
tips.groupby('smoker').apply(top)

# You can pass params to the function using **kwargs.
tips.groupby('smoker').apply(top, n=1, column='total_bill')

# This operation is what describe() does.
display(tips.groupby('smoker')["tip_pct"].describe())

df2 = tips.groupby('smoker')["tip_pct"].apply(lambda x: x.describe())
display(df2)

# Disable the hierarchial indexing.

df = tips.groupby('smoker', group_keys=True).apply(top)
display(df)

df = tips.groupby('smoker', group_keys=False).apply(top)
display(df)

df = pd.DataFrame({
    "data1": np.random.randn(100),
    "data2": np.random.randn(100)
})

display(df.head())

quartiles = pd.cut(df.data1, 4)

print(type(quartiles))

quartiles[:4]

# We can use the series above to groupby.

# We can filter by data2 since we already know the mapping.
grouped = df["data2"].groupby(quartiles)

for k, v in grouped:
    print(k)
    print(v.head(2))

def get_stats(group):
    #print group
    #assert 0
    return pd.Series({
        'min': group.min(),
        'max': group.max(),
        'count': group.count(),
        'mean': group.mean()
    })

df2 = grouped.apply(get_stats)

display(df2)

# Move one level of index to columns.
df2.unstack()

pd.qcut(df.data1, 4).head()

pd.cut(df.data1, 4).head()

s = pd.Series(np.random.randn(6))
s[::2] = np.nan

s

s.fillna(s.mean())

# One can fill nans based on the group.

states = [
    'Ohio', 'New York', 'Vermont', 'Florida', 'Oregon', 'Nevada', 'California',
    'Idaho'
]
group_key = ['East'] * 4 + ['West'] * 4
data = pd.Series(np.random.randn(8), index=states)

data

data[['Vermont', "Nevada", "Idaho"]] = np.nan

data

data.groupby(group_key).mean()

fill_mean = lambda g: g.fillna(g.mean())

data.groupby(group_key).apply(fill_mean)

# Hearts
# Spades
# Clubs
# Diamonds
suits = list("HSCD")
# Values.
card_val = list(list(range(1, 10 + 1)) + [10] * 3)
base_names = ['A'] + list(range(2, 10 + 1)) + list("JQK")
assert len(card_val) == len(base_names)

cards = []
for suit in suits:
    cards.extend(str(num) + suit for num in base_names)

deck = pd.Series(card_val * 4, index=cards)
assert len(deck) == 52
deck.head()

# Draw two cards without replacement.

def draw(deck, n=5):
    return deck.sample(n)


draw(deck)

# Last letter is suit.
get_suit = lambda card: card[-1]

# Draw 2 cards per suit.
# Group by suit, and then get 2 cards from each group.
deck.groupby(get_suit, group_keys=False).apply(draw, n=2)

df = pd.DataFrame({
    'category': ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b'],
    'data': np.random.randn(8),
    'weights': np.random.rand(8)
})

df

grouped = df.groupby('category')

# Aggregation function: dot product between data and weights.
get_wavg = lambda g: np.average(g['data'], weights=g['weights'])

grouped.apply(get_wavg)

close_px = pd.read_csv(
    '~/src/github/pydata-book/examples/stock_px_2.csv',
    parse_dates=True,
    index_col=0)

close_px.head()

close_px.info()

close_px.describe()

# Compute rets.

rets = close_px.pct_change().dropna()

rets.head()

# For each stock compute the correlation with SPX.
spx_corr = lambda x: x.corrwith(x['SPX'])

# Groupby year.
get_year = lambda x: x.year

# For each year, compute the correlation of each stock to SPX.
by_year = rets.groupby(get_year)
by_year.apply(spx_corr)

# For each year, compute the correlation of AAPL and MSFT.
by_year.apply(lambda x: x['AAPL'].corr(x['MSFT']))

import statsmodels.api as sm

def regress(data, yvar, xvars):
    Y = data[yvar]
    X = data[xvars]
    X['intercept'] = 1.0
    result = sm.OLS(Y, X).fit()
    return result.params

by_year.apply(regress, 'AAPL', ['SPX'])

tips.head()

# Aggregate through mean by two indices.
tips.pivot_table(index=['day', 'smoker'])

#?tips.pivot_table

# Compute two metrics by
# - 3 vars: 2 on the index and 1 on the columns
tips.pivot_table(['tip_pct', 'size'],
                 index=['time', 'day'],
                 columns='smoker')

# We can also add summation over each var, so that
# there are values for 2 variables.
tips.pivot_table(['tip_pct', 'size'],
                 index=['time', 'day'],
                 columns='smoker',
                 margins=True)

# You can specify the aggregation function by passing aggfunc.
tips.pivot_table('tip_pct',
                 index=['time', 'smoker'],
                 columns='day',
                 aggfunc=len, margins=True)

pd.crosstab([tips.time, tips.day], tips.smoker, margins=True)q