<reponame>ChrisRichardMiles/crm_m5
# AUTOGENERATED! DO NOT EDIT! File to edit: Training day-by-day models.ipynb (unless otherwise specified).

__all__ = ['main', 'lgb_daily']

# Cell
#export
import os
import logging
import collections
import gc
import json
from itertools import chain
import pandas as pd
import matplotlib.pyplot as plt
import lightgbm as lgb
from lightgbm.callback import record_evaluation
import pickle
from sklearn.model_selection import train_test_split
import neptune.new as neptune
from neptune.new.integrations.lightgbm import NeptuneCallback, create_booster_summary
from fastcore.script import call_parse, Param

from .fe import make_grid_df
from .utils import load_features, dict_to_str
from .metric import WRMSSE

# Cell
def main(path_cfg):
    if type(path_cfg) == str:
        with open(path_cfg, 'r') as f:
            cfg = json.load(f)
    else:
        cfg = path_cfg
    with open(cfg['features_json'], 'r') as f:
        dict_features = json.load(f)

    # if type(cfg['start_test']) == list:
    #     for start_test in cfg['start_test']:
    #         tmp_cfg = cfg.copy()
    #         tmp_cfg['start_test'] = start_test
    #         lgb_daily(tmp_cfg, path_features, cfg['path_data'])

    MODEL_NAME = 'lgb_day_by_day'
    target = 'sales'
    p_horizon = 28
    num_series = 30490
    start_train = cfg['start_train']

    df_stv = pd.read_csv(os.path.join(cfg['path_data_raw'], 'sales_train_evaluation.csv'))
    grid_df, _ = make_grid_df(df_stv)

    ###### Only use items with at least 68 days of sales ######
    first_sale = grid_df[grid_df.sales.notnull()].drop_duplicates('id')
    keep_id = first_sale[(cfg['start_test'] - first_sale.d) >= 68].id.tolist()
    df_stv_trunc = df_stv[df_stv.id.isin(keep_id)]
    grid_df = grid_df[grid_df.id.isin(keep_id)]
    del first_sale
    gc.collect()

    #################### full valid and test sets ###################
    valid_days = [cfg['start_test'] + d - p_horizon for d in range(p_horizon)]
    valid_actuals = df_stv_trunc[[f'd_{d}' for d in valid_days]].values
    e = WRMSSE(cfg['path_data_raw'], cfg['start_test'], df_stv_trunc=df_stv_trunc)
    if cfg['fobj_weight_col'] == 'total_scaled_weight':
        e.add_total_scaled_weight()

    if cfg['start_test'] != 1942: test_actuals = e.actuals.copy()
    prediction_df = df_stv_trunc[['id']]

    del df_stv, df_stv_trunc, e.df_stv
    gc.collect()

    # To avoid problems with the validation and test sets,
    # we fill nan in valid and test period in case we have
    # any nans from fixing out-of-stock data.
    test_days = [cfg['start_test'] + i for i in range(28)]
    mask = grid_df.d.isin(test_days + valid_days) & grid_df[target].isna()
    grid_df.loc[mask, target] = 0

    # Remove data where sales is nan
    grid_df = grid_df[grid_df[target].notnull()]
    gc.collect()

    # Keeping all data to subsample from fro day by day models
    full_grid_df = grid_df.copy()

    # For neptune:
    dict_eval_logs = []
    booster_summaries = []

    ############### Day by day training and predicting #############
    if cfg['days_to_predict'] == "all": cfg['days_to_predict'] = range(28)

    for day_of_horizon in cfg['days_to_predict']:
    # day_of_horizon = 0
        # Starting with full data and filtering for same day of week
        grid_df = full_grid_df.copy()
        test_day = cfg['start_test'] + day_of_horizon
        valid_day = test_day - 28
        same_day_of_week = [d for d in range(cfg['start_train'], test_day + 1) if d%7 == (test_day)%7]
        grid_df = grid_df[grid_df.d.isin(same_day_of_week)]
        gc.collect()

        ######################## Loading Features #######################
        files = [f.replace('csv', 'pkl') for f in dict_features.keys()]
        feature_cols = list(chain(*dict_features.values()))
        if type(cfg['path_features']) == 'str':
            cfg['path_features'] = [cfg['path_features']]
        for folder in cfg['path_features']:
            path = os.path.join('../input', folder)
            for file in os.listdir(path):
                if file in files:
                    tmp = os.path.join(path, file)
                    df = pd.read_pickle(tmp)
                    keep_cols = [c for c in df.columns if c in feature_cols]
                    df = df[keep_cols]
                    if file.startswith('shift'):
                        df.index += day_of_horizon * 30490
                    df = df.reindex(grid_df.index)
                    grid_df = pd.concat([grid_df, df], axis=1)
        del df
        gc.collect()

        remove_features = ['id', 'd', target]
        feature_cols = [col for col in list(grid_df) if col not in remove_features]

        ################## test, train and valid set ####################
        valid_mask = (grid_df.d == valid_day)
        train_mask = (grid_df.d >= start_train) & (grid_df.d < valid_day) & (grid_df[target].notnull())
        test_mask = (grid_df.d == test_day)

        train_x, train_y = grid_df[train_mask][feature_cols], grid_df[train_mask][target]
        valid_x, valid_y = grid_df[valid_mask][feature_cols], grid_df[valid_mask][target]
        test_x, test_y = grid_df[test_mask][feature_cols], grid_df[test_mask][target]

        ################## Fit custom objective and metric ##################
        w_12_train = e.w_12.reindex(grid_df[train_mask].id)
        w_12_eval = e.w_12.reindex(grid_df[valid_mask].id)
        w_12_test = e.w_12.reindex(grid_df[test_mask].id)

        if cfg['fobj']:
            get_fobj = getattr(e, f'get_weighted_{cfg["fobj"]}_fobj')
            fobj = get_fobj(w_12_train, cfg['fobj_weight_col'], cfg['weight_hess'])
        else:
            fobj = None

        if cfg['feval']:
            if cfg['feval'] == 'feval':
                feval = e.feval
            else:
                get_feval = getattr(e, f'get_weighted_{cfg["feval"]}_feval')
                feval = get_feval(w_12_eval, cfg['feval_weight_col'])
        else:
            feval = None

        # Set evaluator actuals to valid day for early stopping
        e.actuals = valid_actuals[:, day_of_horizon].reshape((-1,1))

        ############# lightgbm datasets for training #############
        if cfg['weight_col']:
            weight_train = w_12_train[cfg['weight_col']].values
            weight_eval = w_12_eval[cfg['weight_col']].values
            weight_test = w_12_test[cfg['weight_col']].values
        else:
            weight_train, weight_eval, weight_test = None, None, None
        train_data = lgb.Dataset(train_x, label=train_y, weight=weight_train)
        valid_data = lgb.Dataset(valid_x, label=valid_y, weight=weight_eval, reference = train_data)
#         test_data = lgb.Dataset(test_x, label=test_y, weight=weight_test)

        del grid_df, train_x, train_y, valid_x, valid_y, \
            w_12_train, w_12_eval, w_12_test, \
            weight_train, weight_eval, weight_test
        gc.collect()


        ####################### Training ##########################
        # Callbacks
        dict_eval_log = {}
        # neptune_callback = NeptuneCallback(run=run)

        estimator = lgb.train(
            cfg['lgb_params'],
            train_set=train_data,
            valid_sets=[valid_data],
            valid_names=['valid'],
            fobj = fobj,
            feval = feval,
            callbacks=[record_evaluation(dict_eval_log)],
            verbose_eval=500
        )

        preds = estimator.predict(test_x)
        prediction_df[f'F{day_of_horizon + 1}'] = preds
        gc.collect()

        booster_summary = create_booster_summary(
            booster=estimator,
            max_num_features=50,
            log_trees_as_dataframe=False
        )
        booster_summaries.append(booster_summary)
        dict_eval_logs.append(dict_eval_log)
        del booster_summary, dict_eval_log
        gc.collect()

    # Saving predictions in submission ready format
    tmp = prediction_df.copy()
    prediction_df.id = prediction_df.id.str.replace('evaluation', 'validation')
    #     path = os.path.join(cfg['path_models'], MODEL_NAME + '_' + path_cfg, 'preds.csv')
    #     os.makedirs(path, exist_ok=True)
    path = MODEL_NAME + '_' + path_cfg.split('.')[0] + '_' + 'submission.csv'
    prediction_df = pd.concat([tmp, prediction_df])
    prediction_df.to_csv(path, index=False)

# Cell
@call_parse
def lgb_daily(path_cfg: Param('path to the configuration json', str)='cfg.json'):
    if type(path_cfg) == str:
        with open(path_cfg, 'r') as f:
            cfg = json.load(f)
    else:
        cfg = path_cfg
    with open(cfg['features_json'], 'r') as f:
        dict_features = json.load(f)

    # if type(cfg['start_test']) == list:
    #     for start_test in cfg['start_test']:
    #         tmp_cfg = cfg.copy()
    #         tmp_cfg['start_test'] = start_test
    #         lgb_daily(tmp_cfg, path_features, cfg['path_data'])

    MODEL_NAME = 'lgb_day_by_day'
    target = 'sales'
    p_horizon = 28
    num_series = 30490
    start_train = cfg['start_train']

    df_stv = pd.read_csv(os.path.join(cfg['path_data_raw'], 'sales_train_evaluation.csv'))
    grid_df, _ = make_grid_df(df_stv)

    ###### Only use items with at least 68 days of sales ######
    first_sale = grid_df[grid_df.sales.notnull()].drop_duplicates('id')
    keep_id = first_sale[(cfg['start_test'] - first_sale.d) >= 68].id.tolist()
    df_stv_trunc = df_stv[df_stv.id.isin(keep_id)]
    grid_df = grid_df[grid_df.id.isin(keep_id)]

    #################### full valid and test sets ###################
    valid_days = [cfg['start_test'] + d - p_horizon for d in range(p_horizon)]
    valid_actuals = df_stv_trunc[[f'd_{d}' for d in valid_days]].values
    e = WRMSSE(cfg['path_data_raw'], cfg['start_test'], df_stv_trunc=df_stv_trunc)
    if cfg['fobj_weight_col'] == 'total_scaled_weight':
        e.add_total_scaled_weight()

    if cfg['start_test'] != 1942: test_actuals = e.actuals.copy()
    prediction_df = df_stv_trunc[['id']]

    # Keeping all data to subsample from fro day by day models
    full_grid_df = grid_df.copy()

    # For neptune:
    dict_eval_logs = []
    booster_summaries = []

    ############### Day by day training and predicting #############
    if cfg['days_to_predict'] == "all": cfg['days_to_predict'] = range(28)
    for day_of_horizon in cfg['days_to_predict']:
    #     day_of_horizon = 0

        # Starting with full data and filtering for same day of week
        grid_df = full_grid_df.copy()
        test_day = cfg['start_test'] + day_of_horizon
        valid_day = test_day - 28
        same_day_of_week = [d for d in range(cfg['start_train'], test_day + 1) if d%7 == (test_day)%7]
        grid_df = grid_df[grid_df.d.isin(same_day_of_week)]

        ######################## Loading Features #######################
        index = grid_df.index.tolist()
        grid_df = pd.concat([
            grid_df,
            load_features(cfg['path_features'], dict_features, reindex_with=index, shift_index=num_series * day_of_horizon)
        ], axis=1)

        remove_features = ['id', 'd', target]
        feature_cols = [col for col in list(grid_df) if col not in remove_features]

        ################## test, train and valid set ####################
        valid_mask = (grid_df.d == valid_day)
        train_mask = (grid_df.d >= start_train) & (grid_df.d < valid_day) & (grid_df[target].notnull())
        test_mask = (grid_df.d == test_day)

        train_x, train_y = grid_df[train_mask][feature_cols], grid_df[train_mask][target]
        valid_x, valid_y = grid_df[valid_mask][feature_cols], grid_df[valid_mask][target]
        test_x, test_y = grid_df[test_mask][feature_cols], grid_df[test_mask][target]

        ################## Fit custom objective and metric ##################
        w_12_train = e.w_12.reindex(grid_df[train_mask].id)
        w_12_eval = e.w_12.reindex(grid_df[valid_mask].id)
        w_12_test = e.w_12.reindex(grid_df[test_mask].id)

        if cfg['fobj']:
            get_fobj = getattr(e, f'get_weighted_{cfg["fobj"]}_fobj')
            fobj = get_fobj(w_12_train, cfg['fobj_weight_col'], cfg['weight_hess'])
        else:
            fobj = None

        if cfg['feval']:
            if cfg['feval'] == 'feval':
                feval = e.feval
            else:
                get_feval = getattr(e, f'get_weighted_{cfg["feval"]}_feval')
                feval = get_feval(w_12_eval, cfg['feval_weight_col'])
        else:
            feval = None

        # Set evaluator actuals to valid day for early stopping
        e.actuals = valid_actuals[:, day_of_horizon].reshape((-1,1))

        ############# lightgbm datasets for training #############
        if cfg['weight_col']:
            weight_train = w_12_train[cfg['weight_col']].values
            weight_eval = w_12_eval[cfg['weight_col']].values
            weight_test = w_12_test[cfg['weight_col']].values
        else:
            weight_train, weight_eval, weight_test = None, None, None
        train_data = lgb.Dataset(train_x, label=train_y, weight=weight_train)
        valid_data = lgb.Dataset(valid_x, label=valid_y, weight=weight_eval)
        test_data = lgb.Dataset(test_x, label=test_y, weight=weight_test)


        ####################### Training ##########################
        # Callbacks
        dict_eval_log = {}
        # neptune_callback = NeptuneCallback(run=run)

        estimator = lgb.train(
            cfg['lgb_params'],
            train_set=train_data,
            valid_sets=[valid_data, test_data],
            valid_names=['valid', 'test'],
            fobj = fobj,
            feval = feval,
            callbacks=[record_evaluation(dict_eval_log)],
            verbose_eval=500
        )
        booster_summary = create_booster_summary(booster=estimator, max_num_features=50)
        booster_summary = create_booster_summary(booster=estimator, max_num_features=50)
        booster_summaries.append(booster_summary)
        dict_eval_logs.append(dict_eval_log)

        preds = estimator.predict(test_x)
        prediction_df[f'F{day_of_horizon + 1}'] = preds
        gc.collect()

    # Saving predictions in submission ready format
    tmp = prediction_df.copy()
    prediction_df.id = prediction_df.id.str.replace('evaluation', 'validation')
#     path = os.path.join(cfg['path_models'], MODEL_NAME + '_' + path_cfg, 'preds.csv')
#     os.makedirs(path, exist_ok=True)
    path = MODEL_NAME + '_' + path_cfg.split('.')[0] + '_' + 'submission.csv'
    prediction_df.to_csv(path, index=False)

    ############## Neptune experiment tracking ###############
    if cfg['use_neptune']:
        run = neptune.init(
            project=cfg['neptune_project'],
            api_token=cfg['neptune_api_token'],
            custom_run_id='test',
            name='Untitled',
            description='',
            tags=['testing', '2test'],
            source_files=[cfg['features_json'], path_cfg],
        )
#     else:
#         run = {}

        if cfg['start_test'] != 1942:
            e.actuals = test_actuals
            run['WRMSSE'] = e.score(prediction_df.iloc[:, 1:].values)

        # General params
        for k, v in cfg.items():
            if k not in ['lgb_params', 'neptune_project', 'neptune_api_token']:
                run[k] = v

        # Logs for each days model
        for i, day in enumerate(cfg["days_to_predict"]):
            run[f'lgbm_summary/day_{day}'] = booster_summaries[i]

            for valid_set, odict in dict_eval_log.items():
                for metric, log in odict.items():
                    for val in log:
                        run[f'eval_logs/{day}/{valid_set}/{metric}'].log(val)
        run.stop()
#         if cfg['use_neptune']: run.stop()
#     else:
# #         path = os.path.join(cfg['path_models'], MODEL_NAME + '_' + path_cfg, 'run.json')
# #         os.makedirs(path, exist_ok=True)
#         path = MODEL_NAME + '_' + path_cfg.split('.')[0] + '_' + 'run.json'
#         with open(path, 'w') as f:
#             json.dump(run, f)