import pandas as pd
import numpy as np   
import matplotlib.pyplot as plt
import seaborn as sns
import holidays

df_train = pd.read_csv('../inputs/train.csv', index_col=0, parse_dates=['date'])
df_test = pd.read_csv('../inputs/test.csv', index_col=0, parse_dates=['date'])

df_train.shape, df_test.shape

df_train.head()

def create_date_features(df):
    df['day_of_year'] = df.date.dt.day_of_year
    df['day_of_month'] = df.date.dt.day
    df['day_of_week'] = df.date.dt.weekday
    df['month'] = df.date.dt.month
    df['quarter'] = df.date.dt.quarter
    df['year'] = df.date.dt.year
    df['period'] = df.date.dt.to_period('M')
    return df

df_train = create_date_features(df_train)
df_test = create_date_features(df_test)

df_train.shape, df_test.shape
    

df_train.date.agg({np.min, np.max})

df_test.date.agg({np.min, np.max})

df_train.period.value_counts().sort_index().plot(figsize = (12,8));

df_test.period.value_counts().sort_index().plot(figsize = (12,8));

def cat_per_column(df):
    cat_cols = [col for col in df.columns if df[col].dtype == 'object']
    for col in cat_cols:
        print(df[col].value_counts(), '\n')
        

cat_per_column(df_train)

cat_per_column(df_test)

df_train.set_index('period').num_sold.plot(figsize = (16,8));

df_train.groupby(['country', 'store','product']).num_sold.agg(['min','max','mean'])

df_train.groupby(['country', 'store','product']).num_sold.mean().unstack(level = 'store').assign(ratio = lambda x: x.KaggleRama/x.KaggleMart)

df_train.query('store == "KaggleRama"').groupby(["day_of_week"]).num_sold.mean().plot(label = 'KaggleRama')
df_train.query('store == "KaggleMart"').groupby(["day_of_week"]).num_sold.mean().plot(label = 'KaggleMart')
plt.legend()

df_train.query('product == "Kaggle Sticker"').groupby(["day_of_week"]).num_sold.mean().plot(label = "Kaggle Sticker")
df_train.query('product == "Kaggle Hat"').groupby(["day_of_week"]).num_sold.mean().plot(label = "Kaggle Hat")
df_train.query('product == "Kaggle Mug"').groupby(["day_of_week"]).num_sold.mean().plot(label = "Kaggle Mug")
plt.legend();

df_train.groupby("day_of_month").num_sold.mean().plot();

df_train.query('store == "KaggleRama"').groupby(["day_of_month"]).num_sold.mean().plot(label = 'KaggleRama')
df_train.query('store == "KaggleMart"').groupby(["day_of_month"]).num_sold.mean().plot(label = 'KaggleMart')
plt.legend();

df_train.query('product == "Kaggle Sticker"').groupby(["day_of_month"]).num_sold.mean().plot(label = "Kaggle Sticker")
df_train.query('product == "Kaggle Hat"').groupby(["day_of_month"]).num_sold.mean().plot(label = "Kaggle Hat")
df_train.query('product == "Kaggle Mug"').groupby(["day_of_month"]).num_sold.mean().plot(label = "Kaggle Mug")
plt.legend();

df_train.groupby("day_of_year").num_sold.mean().plot();

df_train.query('store == "KaggleRama"').groupby(["day_of_year"]).num_sold.mean().plot(label = 'KaggleRama')
df_train.query('store == "KaggleMart"').groupby(["day_of_year"]).num_sold.mean().plot(label = 'KaggleMart')
plt.legend();

df_train.query('product == "Kaggle Sticker"').groupby(["day_of_year"]).num_sold.mean().plot(label = 'Kaggle Sticker')
df_train.query('product == "Kaggle Hat"').groupby(["day_of_year"]).num_sold.mean().plot(label = 'Kaggle Hat')
df_train.query('product == "Kaggle Mug"').groupby(["day_of_year"]).num_sold.mean().plot(label = 'Kaggle Mug')
plt.legend();

df_train.groupby("year").num_sold.mean().plot();

df_train.query('store == "KaggleRama"').groupby(["year"]).num_sold.mean().plot(label = 'KaggleRama')
df_train.query('store == "KaggleMart"').groupby(["year"]).num_sold.mean().plot(label = 'KaggleMart')
plt.legend();

df_train.query('product == "Kaggle Sticker"').groupby(["year"]).num_sold.mean().plot(label = 'Kaggle Sticker')
df_train.query('product == "Kaggle Hat"').groupby(["year"]).num_sold.mean().plot(label = 'Kaggle Hat')
df_train.query('product == "Kaggle Mug"').groupby(["year"]).num_sold.mean().plot(label = 'Kaggle Mug')
plt.legend();

plt.figure(figsize = (18,12))
df_train.set_index('date').num_sold.plot(label = 'daily')
df_train.set_index('date').num_sold.resample('W').mean().plot(label = 'weekly')
df_train.set_index('date').num_sold.resample('M').mean().plot(label = 'monthly')
plt.legend();

finland = pd.DataFrame([dict(date = date, 
                            finland_holiday = event, 
                            country= 'Finland') for date, event in holidays.Finland(years=[2015, 2016, 2017, 2018, 2019]).items()])
finland['date'] = finland['date'].astype("datetime64")

norway = pd.DataFrame([dict(date = date, 
                        norway_holiday = event, 
                        country= 'Norway') for date, event in holidays.Norway(years=[2015, 2016, 2017, 2018, 2019]).items()])
norway['date'] = norway['date'].astype("datetime64")

sweden = pd.DataFrame([dict(date = date, 
                            sweden_holiday = event.replace(", Söndag", ""), 
                            country= 'Sweden') for date, event in holidays.Sweden(years=[2015, 2016, 2017, 2018, 2019]).items() if event != 'Söndag'])
sweden['date'] = sweden['date'].astype("datetime64")

df_train = df_train.merge(finland, on = ['date', 'country'], how = 'left').merge(norway, on = ['date', 'country'], how = 'left').merge(sweden, on = ['date', 'country'], how = 'left')
df_test = df_test.merge(finland, on = ['date', 'country'], how = 'left').merge(norway, on = ['date', 'country'], how = 'left').merge(sweden, on = ['date', 'country'], how = 'left')

df_train

from numpy.fft import rfft, irfft, rfftfreq

def low_pass(s, threshold=2e4, d = None):
    if d is None:
        d = 2e-3 / s.size
    fourier = rfft(s)
    frequencies = rfftfreq(s.size, d=d)
    fourier[frequencies > threshold] = 0
    return irfft(fourier)

plt.figure(figsize = (18,12))
s = df_train.num_sold
for i, d in enumerate([1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7, 1e-8, 1e-9], start=1):
    ax = plt.subplot(3,3, i)
    ax.plot(low_pass(s, threshold=2e4, d = d))
    ax.set_title(f'd = {d}')

class CyclicalTransformerV2(CyclicalTransformer):
    def __init__(self, suffix = None, **kwargs):
            super().__init__(**kwargs)
            self.suffix = suffix
        
    def transform(self, X):
        X = super().transform(X)
        if self.suffix is not None:
            transformed_names = X.filter(regex = r'sin$|cos$').columns
            new_names = {name: name + self.suffix for name in transformed_names}
            X.rename(columns=new_names, inplace=True)
        return X

products = pd.get_dummies(df_train['product'])
products.index = df_train.date

fourier_variables = pd.DataFrame({'sin': np.sin(2*np.pi*1*df_train['day_of_week']/7),
                            'cos': np.cos(2*np.pi*1*df_train['day_of_week']/7),
                            'sin2': np.sin(2*np.pi*2*df_train['day_of_week']/7),
                            'cos2': np.cos(2*np.pi*2*df_train['day_of_week']/7),
                            'sin3': np.sin(2*np.pi*3*df_train['day_of_week']/7),
                            'cos3': np.cos(2*np.pi*3*df_train['day_of_week']/7)})

fourier_variables.index = df_train.date
fourier_variables[['sin','cos']].plot(figsize = (16,12));

ax = fourier_variables.plot.scatter('sin', 'cos', figsize = (16,12)).set_aspect('equal')

class Preprocess:
    def __init__(self, gdp_path, easter, special_days, country_holidays):
        self.gdp_path = gdp_path     
        self.easter = easter  
        self.special_days = special_days
        self.country_holidays = country_holidays
    
    def import_data(self, path):
        self.df = pd.read_csv(path, parse_dates=['date'], index_col=0)
        self.index = self.df.index
        
    def get_gdp(self, path):
        gdp = pd.read_csv(path)
        gdp.columns = gdp.columns.str.title().str.replace('Gdp_','')
        gdp = gdp.set_index('Year').stack(0).reset_index()
        gdp.columns = ['year','country','gdp']
        return gdp

    def create_date_features(self, df):
        # A
        df['day_of_year'] = df.date.dt.day_of_year
        df['day_of_month'] = df.date.dt.day
        df['day_of_week'] = df.date.dt.weekday
        df['week_of_year'] = df.date.dt.isocalendar().week.astype('int64')
        df['year'] = df.date.dt.year
        return df
    
    def add_country_holidays(self, df):
        years = [2015, 2016, 2017, 2018, 2019]
        
        finland = self.country_holidays['finland']
        norway = self.country_holidays['norway']
        sweden = self.country_holidays['sweden']
        
        return (df.merge(finland, on = ['date', 'country'], how = 'left')
                        .merge(norway, on = ['date', 'country'], how = 'left')
                        .merge(sweden, on = ['date', 'country'], how = 'left'))

    def get_specific_dates_features(self, df):
        
        return (df.assign(easter = lambda x: x.year.map(self.easter).astype('datetime64'),
                            moms_day = lambda x: x.year.map(self.special_days['moms_day']).astype('datetime64'),
                            wed_jun = lambda x: x.year.map(self.special_days['wed_june']).astype('datetime64'),
                            sun_nov = lambda x: x.year.map(self.special_days['sun_nov']).astype('datetime64'),
                            
                            days_from_easter = lambda x: (x.date - x.easter).dt.days.clip(-5, 65),
                            days_from_mom = lambda x: (x.date - x.moms_day).dt.days.clip(-1, 9),
                            days_from_wed = lambda x: (x.date - x.wed_jun).dt.days.clip(-5, 5),
                            days_from_sun = lambda x: (x.date - x.sun_nov).dt.days.clip(-1, 9),
            )).drop(columns = ['easter','moms_day','wed_jun','sun_nov'])
        
    def join_gdp(self, df):
        return df.merge(self.gdp, on = ['country','year'], how = 'left')

    def feature_engineering(self, df):
        df['log_gdp'] = np.log(df.gdp)
        days_cats = [df.day_of_week < 4, df.day_of_week == 4, df.day_of_week > 4]
        days = ['week','friday','weekends']
        df['week'] = np.select(days_cats, days)
        return df
    
    def __call__(self, path, name = 'Train'):
        self.import_data(path)
        self.gdp = self.get_gdp(self.gdp_path).set_index('year')
        out = (self.df.pipe(self.create_date_features)
                    .pipe(self.join_gdp)
                    .pipe(self.feature_engineering)
                    .pipe(self.add_country_holidays)
                    .pipe(self.get_specific_dates_features))
        out.index = self.index
        print(f'{name} Set created with {out.shape[1]} features')
        return out

from feature_engine.creation import CombineWithReferenceFeature
from feature_engine.encoding import OneHotEncoder, OrdinalEncoder
from feature_engine.imputation import CategoricalImputer
from feature_engine.selection import DropFeatures
from feature_engine.wrappers import SklearnTransformerWrapper
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

from utils import CyclicalTransformerV2

preprocess_dict = {}

scaler = SklearnTransformerWrapper(StandardScaler())
cyc_1 = CyclicalTransformerV2(variables = ['day_of_year'], 
                            max_values = {'day_of_year':365}, 
                            suffix = '_1')
cyc_2 = CyclicalTransformerV2(variables = ['day_of_year'], 
                            max_values = {'day_of_year':365/2,}, 
                            suffix = '_2')
cyc_3 = CyclicalTransformerV2(variables = ['day_of_week'], 
                            max_values = {'day_of_week':7,}, 
                            suffix = '_week')
cyc_4 = CyclicalTransformerV2(variables = ['day_of_week'], 
                            max_values = {'day_of_week':7/2,}, 
                            suffix = '_semiweek')

prep = Pipeline(steps = [
    ('cat_imp', CategoricalImputer()),
    ('ohe', OneHotEncoder(drop_last=True)),
    ('cyc1', cyc_1),
    ('cyc2', cyc_2),
    ('cyc3', cyc_3),
    ('cyc4', cyc_4),
    ('combo', CombineWithReferenceFeature(
        variables_to_combine=['day_of_year_sin_1', 'day_of_year_cos_1',
                           'day_of_year_sin_2', 'day_of_year_cos_2'], 
        reference_variables=['product_Kaggle Mug', 'product_Kaggle Hat'],
        operations = ['mul'])),
    ('combo_week', CombineWithReferenceFeature(
        variables_to_combine=['day_of_week_sin_week', 'day_of_week_cos_week',
        'day_of_week_sin_semiweek', 'day_of_week_cos_semiweek'], 
        reference_variables=['product_Kaggle Mug', 'product_Kaggle Hat'],
        operations = ['mul'])),
    ('drop', DropFeatures(features_to_drop=['year'],
    )),
    ('sc', scaler)
    ])

preprocess_dict['linear_v1'] = prep


bcyc_1 = CyclicalTransformerV2(variables = ['day_of_year'], 
                            max_values = {'day_of_year':365}, 
                            suffix = '_1')
bcyc_2 = CyclicalTransformerV2(variables = ['day_of_year'], 
                            max_values = {'day_of_year':365/2,}, 
                            suffix = '_2')

prep = Pipeline(steps = [
    ('cat_imp', CategoricalImputer()),
    ('oe', OrdinalEncoder()),
    ('cyc1', bcyc_1),
    ('cyc2', bcyc_2),
    ('drop', DropFeatures(features_to_drop=['year']
    )),
])

preprocess_dict['boosting_v1'] = prep

def cv_trainer(model, X,y, X_test = None, folds = None):
    
    S_train = pd.DataFrame(np.nan, index = range(len(X)), columns = ['fold','num_sold'])
    S_test = {}
    
    scores = dict(train = [],
                val = [])
    
    for fold, (train_idx, val_idx) in enumerate(folds):
        
        X_train, y_train = X.iloc[train_idx],y.iloc[train_idx]
        X_val, y_val = X.iloc[val_idx],y.iloc[val_idx]
        
        id_train = X_train.index
        id_val = X_val.index
        id_test = X_test.index
        
        model.fit(X_train, y_train)
        y_pred_train = pd.Series(model.predict(X_train), index = id_train, name = 'num_sold')
        y_pred_val = pd.Series(model.predict(X_val), index = id_val, name = 'num_sold')
        y_pred_test = pd.Series(model.predict(X_test), index = id_test, name = 'num_sold')
        
        S_train.loc[val_idx, 'num_sold'] = y_pred_val
        S_train.loc[val_idx, 'fold'] = fold
        
        scores['train'].append(smape(y_train, y_pred_train))
        scores['val'].append(smape(y_val, y_pred_val))
        
        S_test[fold] = y_pred_test
        
    S_test = pd.DataFrame(S_test).mean(axis = 1).to_frame()
    S_test.columns = ['num_sold']
    
    return scores, S_train, S_test

class HybridModel(BaseEstimator, RegressorMixin):
    def __init__(self, l_model, b_model):
        self.l_model = l_model
        self.b_model = b_model
        
    def fit(self, X, y):
        
        log_y = np.log(y)
        self.l_model.fit(X, log_y)
        linear_pred = np.exp(self.l_model.predict(X))
        y_resid = y - linear_pred
        
        self.b_model.fit(X, y_resid)
        
        self.is_fitted_ = True
        return self
    
    def predict(self, X):
        check_is_fitted(self, 'is_fitted_')
        l_predict = np.exp(self.l_model.predict(X))
        b_predict = self.b_model.predict(X)
        
        y_predict = l_predict + b_predict
        return y_predict

import os

import numpy as np
import pandas as pd
from sklearn.linear_model import HuberRegressor, Ridge, LinearRegression
from sklearn.model_selection import GroupKFold
from sklearn.pipeline import Pipeline
from lightgbm import LGBMRegressor

import data as d
from models import HybridModel
from preprocessing import preprocess_dict
from utils import Preprocess, cv_trainer

GDP_PATH = '../inputs/GDP_data_2015_to_2019_Finland_Norway_Sweden.csv'
TRAIN_PATH = '../inputs/train.csv'
TEST_PATH = '../inputs/test.csv'
feature_engineer = Preprocess(GDP_PATH, d.EASTER_DICT, d.SPECIAL_DAYS_DICT, d.COUNTRY_HOLIDAYS_DICT)
train_df = feature_engineer(TRAIN_PATH)
test_df = feature_engineer(TEST_PATH, name = 'Test')

X = train_df.drop(columns = ['gdp','date','num_sold'])
y = train_df.num_sold
X_test = test_df.drop(columns = ['gdp','date'])

ridge = Pipeline(steps = [
    ('prep', preprocess_dict['linear_v1']),
    ('model', LinearRegression())
])

xgb = Pipeline(steps = [
    ('prep', preprocess_dict['boosting_v1']),
    ('model', LGBMRegressor())
])

model = HybridModel(ridge, xgb)
folds = GroupKFold(n_splits=4).split(X = train_df, groups=train_df.year) # not sure if it's ok
scores, S_train, S_test = cv_trainer(model, X,y, X_test, folds)

model_name = '_'.join([type(model).__name__, type(model.l_model.named_steps.model).__name__, type(model.b_model.named_steps.model).__name__])
print(f'Training results for {model_name}')
print(scores)
print('Mean Training Score: ', np.mean(scores['train']))
print('Mean Validaton Score: ', np.mean(scores['val']))

S_test.to_csv(f'../submissions/submission_{model_name}_v2.csv')

def save_predictions(preds, model_name, mode = 'train'):
    path = f'../submissions/Stacking_{mode}.csv'
    if os.path.exists(path):
        input = pd.read_csv(path, index_col=0)
        input[f'{model_name}'] = preds['num_sold'].sort_index()
        input.to_csv(path)
    else:
        out = preds['num_sold']
        out.name = model_name
        out.to_csv(path)

save_predictions(S_train, model_name)
save_predictions(S_test, model_name, mode = 'test')