# Setup feedback system
from learntools.core import binder
binder.bind(globals())
from learntools.time_series.ex6 import *

# Setup notebook
from pathlib import Path
import ipywidgets as widgets
from learntools.time_series.style import *  # plot style settings
from learntools.time_series.utils import (create_multistep_example,
                                          load_multistep_data,
                                          make_lags,
                                          make_multistep_target,
                                          plot_multistep)

import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.multioutput import RegressorChain
from sklearn.preprocessing import LabelEncoder
from xgboost import XGBRegressor


comp_dir = Path('../input/store-sales-time-series-forecasting')

store_sales = pd.read_csv(
    comp_dir / 'train.csv',
    usecols=['store_nbr', 'family', 'date', 'sales', 'onpromotion'],
    dtype={
        'store_nbr': 'category',
        'family': 'category',
        'sales': 'float32',
        'onpromotion': 'uint32',
    },
    parse_dates=['date'],
    infer_datetime_format=True,
)
store_sales['date'] = store_sales.date.dt.to_period('D')
store_sales = store_sales.set_index(['store_nbr', 'family', 'date']).sort_index()

family_sales = (
    store_sales
    .groupby(['family', 'date'])
    .mean()
    .unstack('family')
    .loc['2017']
)

test = pd.read_csv(
    comp_dir / 'test.csv',
    dtype={
        'store_nbr': 'category',
        'family': 'category',
        'onpromotion': 'uint32',
    },
    parse_dates=['date'],
    infer_datetime_format=True,
)
test['date'] = test.date.dt.to_period('D')
test = test.set_index(['store_nbr', 'family', 'date']).sort_index()

"""
Explanation: Introduction
Run this cell to set everything up!
End of explanation
"""


datasets = load_multistep_data()

data_tabs = widgets.Tab([widgets.Output() for _ in enumerate(datasets)])
for i, df in enumerate(datasets):
    data_tabs.set_title(i, f'Dataset {i+1}')
    with data_tabs.children[i]:
        display(df)

display(data_tabs)

"""
Explanation: Consider the following three forecasting tasks:
a. 3-step forecast using 4 lag features with a 2-step lead time<br>
b. 1-step forecast using 3 lag features with a 1-step lead time<br>
c. 3-step forecast using 4 lag features with a 1-step lead time<br>
Run the next cell to see three datasets, each representing one of the tasks above.
End of explanation
"""


# YOUR CODE HERE: Match the task to the dataset. Answer 1, 2, or 3.
task_a = ____
task_b = ____
task_c = ____

# Check your answer
q_1.check()

# Lines below will give you a hint or solution code
#_COMMENT_IF(PROD)_
q_1.hint()
#_COMMENT_IF(PROD)_
q_1.solution()

#%%RM_IF(PROD)%%
task_a = 2
task_b = 1
task_c = 3

q_1.assert_check_passed()

"""
Explanation: 1) Match description to dataset
Can you match each task to the appropriate dataset?
End of explanation
"""


print("Training Data", "\n" + "-" * 13 + "\n", store_sales)
print("\n")
print("Test Data", "\n" + "-" * 9 + "\n", test)

"""
Explanation: Look at the time indexes of the training and test sets. From this information, can you identify the forecasting task for Store Sales?
End of explanation
"""


# View the solution (Run this cell to receive credit!)
q_2.check()

"""
Explanation: 2) Identify the forecasting task for Store Sales competition
Try to identify the forecast origin and the forecast horizon. How many steps are within the forecast horizon? What is the lead time for the forecast?
Run this cell after you've thought about your answer.
End of explanation
"""


# YOUR CODE HERE
y = family_sales.loc[:, 'sales']

# YOUR CODE HERE: Make 4 lag features
X = ____

# YOUR CODE HERE: Make multistep target
y = ____

#_UNCOMMENT_IF(PROD)_
#y, X = y.align(X, join='inner', axis=0)

# Check your answer
q_3.check()

# Lines below will give you a hint or solution code
#_COMMENT_IF(PROD)_
q_3.hint()
#_COMMENT_IF(PROD)_
q_3.solution()

#%%RM_IF(PROD)%%
y = family_sales.loc[:, 'sales']
X = make_lags(y, lags=4).dropna()
y = make_multistep_target(y, steps=16).dropna()
y, X = y.align(X, join='inner', axis=0)

q_3.assert_check_passed()

"""
Explanation: In the tutorial we saw how to create a multistep dataset for a single time series. Fortunately, we can use exactly the same procedure for datasets of multiple series.
3) Create multistep dataset for Store Sales
Create targets suitable for the Store Sales forecasting task. Use 4 days of lag features. Drop any missing values from both targets and features.
End of explanation
"""


le = LabelEncoder()
X = (X
    .stack('family')  # wide to long
    .reset_index('family')  # convert index to column
    .assign(family=lambda x: le.fit_transform(x.family))  # label encode
)
y = y.stack('family')  # wide to long

display(y)

"""
Explanation: In the tutorial, we saw how to forecast with the MultiOutput and Direct strategies on the Flu Trends series. Now, you'll apply the DirRec strategy to the multiple time series of Store Sales.
Make sure you've successfully completed the previous exercise and then run this cell to prepare the data for XGBoost.
End of explanation
"""


from sklearn.multioutput import RegressorChain

# YOUR CODE HERE
model = ____

# Check your answer
q_4.check()

# Lines below will give you a hint or solution code
#_COMMENT_IF(PROD)_
q_4.hint()
#_COMMENT_IF(PROD)_
q_4.solution()

#%%RM_IF(PROD)%%
from sklearn.multioutput import RegressorChain

model = RegressorChain(XGBRegressor())
q_4.assert_check_passed()

"""
Explanation: 4) Forecast with the DirRec strategy
Instatiate a model that applies the DirRec strategy to XGBoost.
End of explanation
"""


model.fit(X, y)

y_pred = pd.DataFrame(
    model.predict(X),
    index=y.index,
    columns=y.columns,
).clip(0.0)

"""
Explanation: Run this cell if you'd like to train this model.
End of explanation
"""


FAMILY = 'BEAUTY'
START = '2017-04-01'
EVERY = 16

y_pred_ = y_pred.xs(FAMILY, level='family', axis=0).loc[START:]
y_ = family_sales.loc[START:, 'sales'].loc[:, FAMILY]

fig, ax = plt.subplots(1, 1, figsize=(11, 4))
ax = y_.plot(**plot_params, ax=ax, alpha=0.5)
ax = plot_multistep(y_pred_, ax=ax, every=EVERY)
_ = ax.legend([FAMILY, FAMILY + ' Forecast'])

"""
Explanation: And use this code to see a sample of the 16-step predictions this model makes on the training data.
End of explanation
"""


