import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from itertools import product
from sklearn.preprocessing import LabelEncoder



# =============================================================================
# The lines where we processed our data
# =============================================================================
def lag_feature(df, lags, col):
    tmp = df[['date_block_num','shop_id','item_id',col]]
    for i in lags:
        shifted = tmp.copy()
        shifted.columns = ['date_block_num','shop_id','item_id', col+'_lag_'+str(i)]
        shifted['date_block_num'] += i
        df = pd.merge(df, shifted, on=['date_block_num','shop_id','item_id'], how='left')
    return df



items = pd.read_csv(r'dataset\items.csv')
shops = pd.read_csv(r'dataset\shops.csv')
cats = pd.read_csv(r'dataset\item_categories.csv')
train = pd.read_csv(r'dataset\sales_train.csv')
test  = pd.read_csv(r'dataset\test.csv').set_index('ID')

train = train[train.item_price<100000]
train = train[train.item_cnt_day<1001]

median = train[(train.shop_id==32)&(train.item_id==2973)&(train.date_block_num==4)&(train.item_price>0)].item_price.median()
train.loc[train.item_price<0, 'item_price'] = median

train.loc[train.shop_id == 0, 'shop_id'] = 57
test.loc[test.shop_id == 0, 'shop_id'] = 57
train.loc[train.shop_id == 1, 'shop_id'] = 58
test.loc[test.shop_id == 1, 'shop_id'] = 58
train.loc[train.shop_id == 10, 'shop_id'] = 11
test.loc[test.shop_id == 10, 'shop_id'] = 11


shops['shop_name'] = shops['shop_name'].apply(lambda x: x.lower()).str.replace('[^\w\s]', '').str.replace('\d+','').str.strip()
shops['city'] = shops['shop_name'].str.partition(' ')[0]
shops['city_code'] = LabelEncoder().fit_transform(shops['city'])
shops['shop_type'] = shops['shop_name'].apply(lambda x: 'мтрц' if 'мтрц' in x else 'трц' if 'трц' in x else 'трк' if 'трк' in x else 'тц' if 'тц' in x else 'тк' if 'тк' in x else 'NO_DATA')
shops['shop_type'] = LabelEncoder().fit_transform(shops['shop_type'])
shops = shops[['shop_id','city_code','shop_type']]


cats['split'] = cats['item_category_name'].str.split('-')
cats['type'] = cats['split'].map(lambda x: x[0].strip())
cats['type_code'] = LabelEncoder().fit_transform(cats['type'])
# if subtype is nan then type
cats['subtype'] = cats['split'].map(lambda x: x[1].strip() if len(x) > 1 else x[0].strip())
cats['subtype_code'] = LabelEncoder().fit_transform(cats['subtype'])
cats = cats[['item_category_id','type_code', 'subtype_code']]
items.drop(['item_name'], axis=1, inplace=True)


matrix = []
cols = ['date_block_num','shop_id','item_id']
for i in range(34):
    sales = train[train.date_block_num==i]
    matrix.append(np.array(list(product([i], sales.shop_id.unique(), sales.item_id.unique())), dtype='int16'))
    
matrix = pd.DataFrame(np.vstack(matrix), columns=cols)
matrix['date_block_num'] = matrix['date_block_num'].astype(np.int8)
matrix['shop_id'] = matrix['shop_id'].astype(np.int8)
matrix['item_id'] = matrix['item_id'].astype(np.int16)
matrix.sort_values(cols,inplace=True)


train['revenue'] = train['item_price'] *  train['item_cnt_day']

item_price_lag = train.groupby(['date_block_num','item_id']).agg({'item_price':['mean']})
item_price_lag.columns = ['average_item_price']
item_price_by_shop_lag = train.groupby(['date_block_num','shop_id', 'item_id']).agg({'item_price':['mean']})
item_price_by_shop_lag.columns = ['average_item_price_by_shop']
group = train.groupby(['date_block_num','shop_id','item_id']).agg({'item_cnt_day': ['sum']})
group.columns = ['item_cnt_month']
group.reset_index(inplace=True)

matrix = pd.merge(matrix, group, on=cols, how='left')
matrix['item_cnt_month'] = (matrix['item_cnt_month'].fillna(0).clip(0,20).astype(np.float16))



test['date_block_num'] = 34
test['date_block_num'] = test['date_block_num'].astype(np.int8)
test['shop_id'] = test['shop_id'].astype(np.int8)
test['item_id'] = test['item_id'].astype(np.int16)


matrix = pd.concat([matrix, test], ignore_index=True, sort=False, keys=cols)
matrix.fillna(0, inplace=True) # 34 month

matrix = pd.merge(matrix, item_price_lag, on=['date_block_num','item_id'], how='left')
matrix['average_item_price'] = matrix['average_item_price'].astype(np.float16)
matrix = lag_feature(matrix, [1,2,3], 'average_item_price')
matrix.drop(['average_item_price'], axis=1, inplace=True)
matrix = pd.merge(matrix, item_price_by_shop_lag, on=['date_block_num','shop_id','item_id'], how='left')
matrix['average_item_price_by_shop'] = matrix['average_item_price_by_shop'].astype(np.float16)
matrix = lag_feature(matrix, [1,2,3], 'average_item_price_by_shop')
matrix.drop(['average_item_price_by_shop'], axis=1, inplace=True)



matrix = pd.merge(matrix, shops, on=['shop_id'], how='left')
matrix = pd.merge(matrix, items, on=['item_id'], how='left')
matrix = pd.merge(matrix, cats, on=['item_category_id'], how='left')
matrix['city_code'] = matrix['city_code'].astype(np.int8)
matrix['shop_type'] = matrix['shop_type'].astype(np.int8)
matrix['item_category_id'] = matrix['item_category_id'].astype(np.int8)
matrix['type_code'] = matrix['type_code'].astype(np.int8)
matrix['subtype_code'] = matrix['subtype_code'].astype(np.int8)



shop_mean = matrix.groupby(['shop_id']).agg({'item_cnt_month': ['mean']})
shop_mean.columns = ['shop_mean']
shop_mean.reset_index(inplace=True)
shop_item_mean = matrix.groupby(['item_id','shop_id']).agg({'item_cnt_month': ['mean']})
shop_item_mean.columns = ['shop_item_mean']
shop_item_mean.reset_index(inplace=True)
group = matrix.groupby(['date_block_num', 'item_id']).agg({'item_cnt_month': ['mean']})
group.columns = [ 'date_item_avg_item_cnt' ]
group.reset_index(inplace=True)

matrix = pd.merge(matrix, shop_mean, on=['shop_id'], how='left')
matrix = pd.merge(matrix, shop_item_mean, on=['item_id','shop_id'], how='left')
matrix = pd.merge(matrix, group, on=['date_block_num','item_id'], how='left')
matrix['date_item_avg_item_cnt'] = matrix['date_item_avg_item_cnt'].astype(np.float16)
matrix = lag_feature(matrix, [1,2,3], 'date_item_avg_item_cnt')
matrix.drop(['date_item_avg_item_cnt'], axis=1, inplace=True)
matrix = lag_feature(matrix, [1,2,3], 'item_cnt_month')

matrix_last = matrix[matrix.date_block_num > 2]


def fill_na(df):
    for col in df.columns:
        if ('_lag_' in col) & (df[col].isnull().any()):
            if ('item_cnt' in col):
                df[col].fillna(0, inplace=True)   
            if ('shop_mean' in col):
                df[col].fillna(0, inplace=True)
            if ('average_item_price' in col):
                df[col].fillna(0, inplace=True)
    return df

matrix = fill_na(matrix_last)

matrix_last.to_pickle('dataset/traintest.pkl')


# =============================================================================
# correlation Matrix
# =============================================================================
cor_data = matrix_last[['shop_item_mean','date_block_num','date_item_avg_item_cnt_lag_1','item_category_id','average_item_price_lag_2','average_item_price_lag_1','item_cnt_month_lag_1','item_cnt_month']]
corr = cor_data.corr()
mask = np.zeros_like(corr, dtype=np.bool)

f,ax = plt.subplots(figsize=(15, 20))
cmap = sns.diverging_palette(220, 10, as_cmap=True)
sns.heatmap(corr, mask=mask, cmap=cmap, vmax=.3, center=0,
            square=True, linewidths=.5, cbar_kws={"shrink": .5},annot=True)
plt.savefig('outputdata/correlation.png')
