import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('seaborn')
plt.rcParams['figure.figsize'] = [12, 6]

from hcrystalball.utils import get_sales_data

df = get_sales_data(n_dates=100, 
                    n_assortments=2, 
                    n_states=2, 
                    n_stores=2)

df

from hcrystalball.model_selection import ModelSelector

ms = ModelSelector(horizon=10, 
                   frequency='D', 
                   country_code_column='HolidayCode',
                  )

ms.create_gridsearch(sklearn_models=True,
                    n_splits = 2,
                    between_split_lag=None,
                    sklearn_models_optimize_for_horizon=False,
                    autosarimax_models=False,
                    prophet_models=False,
                    tbats_models=False,
                    exp_smooth_models=False,
                    average_ensembles=False,
                    stacking_ensembles=False,                    
                    exog_cols=['Open','Promo','SchoolHoliday','Promo2'],
                    )

from hcrystalball.wrappers import get_sklearn_wrapper    
from sklearn.linear_model import LinearRegression

ms.add_model_to_gridsearch(get_sklearn_wrapper(LinearRegression))

# from prefect.engine.executors import LocalDaskExecutor
ms.select_model(df=df,
                target_col_name='Sales',
                partition_columns=['Assortment', 'State','Store'],
#                 parallel_over_columns=['Assortment'],
#                 persist_model_selector_results=False,
#                 output_path='my_results',
#                 executor = LocalDaskExecutor(),                
               )

ms.plot_best_wrapper_classes()

ms.plot_results(plot_from='2015-06-01')

ms.results[0]

res = ms.get_result_for_partition(partition=ms.results[0].partition)

res.plot_result(plot_from = '2015-06-01', title='forecasts');

res.plot_error(title='Errors');

from hcrystalball.model_selection import load_model_selector
from hcrystalball.model_selection import load_model_selector_result
from hcrystalball.model_selection import load_best_model

res.persist(path='tmp')
res = load_model_selector_result(path='tmp',partition_label=ms.results[0].partition) 

ms.persist_results(folder_path='tmp/results')
ms = load_model_selector(folder_path='tmp/results')

ms = load_model_selector(folder_path='tmp/results')

ms.plot_results(plot_from='2015-06-01')

# cleanup
import shutil
try:
    shutil.rmtree('tmp')
except:
    pass