# Loading Dataset
import pandas as pd
from ads.dataset.factory import DatasetFactory
from ads.dataset.dataset import ADSDataset

# Exploratory Data Analysis and Data Cleaning
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from datetime import timedelta
from datetime import datetime
from datetime import date

# Model Training
import xgboost as xgb
import time
import gc

# Model Evaluation
from sklearn.metrics import mean_squared_error as MSE

# Hyperparameter Tuning
from sklearn.model_selection import TimeSeriesSplit
from skopt import BayesSearchCV
from sklearn.metrics import make_scorer


# Loading CSV Locally
def read_data_csv(filename, target, targetfillna=True):
    """Load data from csv file"""
    dataset = pd.read_csv(filename)
    if targetfillna:
        dataset.fillna({target: 0}, inplace=True)
    print("Dataset Shape: ", dataset.shape)
    return dataset


# Converting dataframe into OCI Dataset form
def convert_datasetfactory(df, target = "SHIPMENTS_HISTORY"):
    """Convert a Pandas Dataframe into ADS DatasetFactory type"""
    return DatasetFactory.from_dataframe(df, target=target, type_discovery=False)


# Sort Data
def sort_data(df, by, ascending_arr, reset_index):
    """Sort data by a provided list of field columns"""
    sorted_df = df.sort_values(by=by, ascending=ascending_arr)
    if reset_index:
        sorted_df.reset_index(drop=True, inplace=True)
    return sorted_df


# Convert fields
def convert_field_types(df, field_type_dict):
    """Convert the field types of the provided field names in a dataset"""
    return (df.astype(field_type_dict))


# Customer Site lists of Item Numbers, Item Classes, and Customer Sites to be referenced
def get_unique_field_values(df, field_name):
    """Return the list of unique values in a field column"""
    unique_list = list(df[field_name].unique())
    unique_list.sort()
    return unique_list


# Find the first occurrence and last occurrence of each item per customer site
def fill_rows_missing_dates(df, date_field, target_field, item_level, org_level=""):
    """Fill in 0 shipments for days missing between the first occurrence and last occurrence of each item-customer pair"""
    # Find the first occurrence and last occurrence of each item per Customer Site
    if org_level != "":
        forecasting_lvl = [org_level, item_level]
    else:
        forecasting_lvl = [item_level]

    min_date_item = df.df.groupby(forecasting_lvl)[
        date_field].min().reset_index()
    max_date_item = df.df.groupby(forecasting_lvl)[
        date_field].max().reset_index()

    # Removing Items that have no sales records at a CUSTOMER_SITE
    min_date_item.dropna(inplace=True)
    max_date_item.dropna(inplace=True)
    min_date_item = min_date_item.rename(columns={date_field: "Min_DATE"})
    max_date_item = max_date_item.rename(columns={date_field: "Max_DATE"})
    item_date_range = min_date_item.merge(max_date_item, on=forecasting_lvl)

    # Building a dataframe with records for days when each item had no sales recorded
    date_extended = pd.DataFrame()
    if org_level != "":
        for ind in item_date_range.index:
            date_extended_i = df.df[(df.df[item_level] == item_date_range[item_level][ind]) & (
                df.df[org_level] == item_date_range[org_level][ind])][[item_level, org_level, target_field, date_field]].set_index(date_field)
            date_extended_i = date_extended_i.reindex(pd.date_range(
                item_date_range["Min_DATE"][ind], item_date_range["Max_DATE"][ind]), fill_value=np.NaN)
            date_extended_i[item_level].fillna(
                item_date_range[item_level][ind], inplace=True)
            date_extended_i[target_field].fillna(0, inplace=True)
            date_extended_i[org_level].fillna(
                item_date_range[org_level][ind], inplace=True)
            date_extended = pd.concat([date_extended, date_extended_i], axis=0)
        date_extended.index.name = date_field
        date_extended.reset_index(inplace=True)
    else:
        for ind in item_date_range.index:
            date_extended_i = df.df[(df.df[item_level] == item_date_range[item_level][ind])][[
                item_level, target_field, date_field]].set_index(date_field)
            date_extended_i = date_extended_i.reindex(pd.date_range(
                item_date_range["Min_DATE"][ind], item_date_range["Max_DATE"][ind]), fill_value=np.NaN)
            date_extended_i[item_level].fillna(
                item_date_range[item_level][ind], inplace=True)
            date_extended_i[target_field].fillna(0, inplace=True)
            date_extended = pd.concat([date_extended, date_extended_i], axis=0)
        date_extended.index.name = date_field
        date_extended.reset_index(inplace=True)

    return DatasetFactory.from_dataframe(date_extended, target=target_field, type_discovery=False)


# Fill in data for dates when Customer Site was closed
def fill_closure_dates(df, date_field, day_related_field_1, day_related_field_2):
    """Fill in missing data on holidays/Customer Site closure days"""
    date_features = df.df[[date_field, day_related_field_1, day_related_field_2]].groupby(
        date_field).first().reset_index()

    date_features['ORIG_' +
                  day_related_field_1] = date_features.loc[:, day_related_field_1]
    date_features = date_features.set_index(date_field)

    first_date = df.df[[date_field, day_related_field_1, day_related_field_2]].groupby(
        date_field).first().reset_index().iloc[0][date_field]
    last_date = df.df[[date_field, day_related_field_1, day_related_field_2]].groupby(
        date_field).first().reset_index().iloc[-1][date_field]

    idx = pd.date_range(start=first_date, end=last_date, freq="D")
    date_features = date_features.reindex(idx, fill_value=np.nan)

    date_features = date_features.rename_axis(date_field).reset_index()
    xmas_idx = date_features[date_features["HOLIDAY_INDICATOR_1"].isna() & (
        date_features["DATE"].dt.strftime('%m-%d') == '12-25')]["HOLIDAY_INDICATOR_1"].index
    nyd_idx = date_features[date_features["HOLIDAY_INDICATOR_1"].isna() & (
        date_features["DATE"].dt.strftime('%m-%d') == '01-01')]["HOLIDAY_INDICATOR_1"].index

    for i in xmas_idx:
        date_features.at[i, "HOLIDAY_INDICATOR_1"] = 1
    for i in nyd_idx:
        date_features.at[i, "HOLIDAY_INDICATOR_1"] = 1

    return date_features


# Calculating Datetime related fields:
def week_of_month(dt):
    """Returns the week of the month for the specified date."""
    week_of_month_ls = []
    for date in dt:
        day = date.day
        if day < 8:
            week_of_month_ls.append(1)
        elif day < 15:
            week_of_month_ls.append(2)
        elif day < 22:
            week_of_month_ls.append(3)
        elif day < 29:
            week_of_month_ls.append(4)
        elif day < 32:
            week_of_month_ls.append(5)
        else:
            week_of_month_ls.append(np.nan)
    return week_of_month_ls


# Mark weekend dates
def is_weekend(dt):
    """Returns 1 if date is a Saturday or Sunday, and 0 otherwise"""
    weekend_ls = []
    for date in dt:
        if date.strftime("%A") == "Sunday" or date.strftime("%A") == "Saturday":
            weekend_ls.append(1)
        else:
            weekend_ls.append(0)
    return weekend_ls


# Compute calculated date fields
def calc_date_field(dataset, date_col, new_col_name, new_col_code):
    """Adding a new datetime field to an existing dataset"""
    if new_col_code == "week_in_month":
        return dataset.assign_column(new_col_name, week_of_month(dataset[date_col]))
    elif new_col_code == "weekend":
        return dataset.assign_column(new_col_name, is_weekend(dataset[date_col]))
    elif new_col_code[0] == "%":
        return dataset.assign_column(new_col_name, (dataset.df[date_col].dt.strftime(new_col_code).astype(int)))
    elif new_col_code in dataset.columns:
        # field name already exists
        return dataset
    else:
        return dataset


# Call calc_date_field()
def calc_all_date_related_fields(date_df, date_field, new_calc_field_dict):
    """Compute all date related fields for a given dataset"""
    if isinstance(date_df, pd.DataFrame):
        date_df = DatasetFactory.from_dataframe(date_df)
    for field_name, field_format in new_calc_field_dict.items():
        date_df = calc_date_field(
            date_df, date_field, field_name, field_format)
    return date_df


# Filling Missing values in CUSTOM_MEASURE_1 with forward fill
def ffill_field(df, field):
    """Forward fill missing values in a given field column"""
    df.df[field] = df.df[field].fillna(method="ffill")
    na_ls = df[df[field].isna()].index
    for i in na_ls:
        df.df[field][i] = df.df[field][na_ls[-1]+1]
    return df


# Data splitting and data joining at different hierarchies to update with new rows of 0 unit sale records
def data_splitting(df):
    """Split dataset into corresponding sub-datasets by hierarchy types"""
    item_features_df = df.df[["CATEGORY_LEVEL_2", "CATEGORY_LEVEL_1",
                              "ITEM", "LIFECYCLE_PHASE"]].groupby("ITEM").first().reset_index()
    group_holiday_ind_2_df = df.df.groupby(["DATE", "CUSTOMER_SITE"]).sum(numeric_only=True)[df.df.groupby(
        ["DATE", "CUSTOMER_SITE"]).sum(numeric_only=True)["HOLIDAY_INDICATOR_2"] > 0].reset_index()[["DATE", "CUSTOMER_SITE"]]
    group_holiday_ind_2_df["HOLIDAY_INDICATOR_2"] = 1
    promo_df = df.df[["DATE", "ITEM", "CUSTOMER_SITE", "EVENT_TYPE"]]
    customer_site_features_df = df.df[["CUSTOMER_SITE", "CUSTOMER", "ZONE"]]
    customer_site_features_df = customer_site_features_df.drop_duplicates(
    ).sort_values(by=["CUSTOMER_SITE"])
    return item_features_df, group_holiday_ind_2_df, promo_df, customer_site_features_df


def join_all_data(df, date_features_df, item_features_df, group_holiday_ind_2_df, promo_df, customer_site_features_df):
    """Merge all sub-datasets into a complete dataset"""
    df = df.merge(date_features_df, how="left", on="DATE")
    df = df.merge(group_holiday_ind_2_df, how="left",
                  on=["DATE", "CUSTOMER_SITE"])
    df.df["HOLIDAY_INDICATOR_2"] = df.df["HOLIDAY_INDICATOR_2"].fillna(0)
    df = df.merge(item_features_df, how="left", on="ITEM")
    df = df.merge(promo_df, how="left", on=["DATE", "ITEM", "CUSTOMER_SITE"])
    df.df["EVENT_TYPE"] = df.df["EVENT_TYPE"].fillna(0)
    df = df.merge(customer_site_features_df, how="left", on=["CUSTOMER_SITE"])
    return df


# Fill in data for dates when Customer Site was closed for an extended period
def handle_customer_site_closure(df, customer_site_field, customer_site_label, date_field, closure_start_date, closure_end_date):
    """Fill in rows of data for long-term customer site closures"""
    df_without_sales_during_closure = df.df[(df[customer_site_field] == customer_site_label) & (
        (df.df[date_field] < closure_start_date) | (df.df[date_field] > closure_end_date))]
    return pd.concat([df.df[df.df[customer_site_field] != customer_site_label], df_without_sales_during_closure])


# Partition by Customer Site
# Use cleaned dataset for partitioning and training
def partition_by_customer_site(df, partition_level, partition_level_str_ls):
    """Split the data into separate datasets based on selected parition level"""
    partition = [x for _, x in df.groupby(df[partition_level])]
    dict_data = {}
    for i, name in enumerate(partition_level_str_ls):
        dict_data[name] = {"full": partition[i]}
    del partition
    return dict_data


# Separate data into training, validation, and testing set
def create_data_split(partition, data, min_date=None, max_date=None):
    """Create a data split from min_date (inclusive) to max_date (exclusive)"""
    if min_date == None and max_date == None:
        return partition[data]
    elif min_date == None:
        data_split = partition[data][partition[data]["DATE"] < max_date]
    elif max_date == None:
        data_split = partition[data][partition[data]["DATE"] >= min_date]
    else:
        data_split = partition[data][(partition[data]["DATE"] >= min_date) & (
            partition[data]["DATE"] < max_date)]
    data_split.reset_index(drop=True, inplace=True)
    return data_split


# Encoding Categorical Fields
# One hot encode field
def one_hot_categorical(data, field_name):
    """One hot encodes provided categorical field"""
    return pd.get_dummies(data[field_name], prefix="is", dtype=int)


# Renaming one hot encoded fields
def rename_one_hot_encoded_fields(one_hot_df, curr_field_names, one_hot_encoded_field):
    """Rename one hot encoded fields with their respective category level names"""
    month_names = np.array(["January", "February", "March", "April", "May",
                           "June", "July", "August", "September", "October", "November", "December"])
    week_names = np.array(["First_Week", "Second_Week",
                          "Third_Week", "Fourth_Week", "Fifth_Week"])
    new_name = "is_"
    for i in np.arange(len(curr_field_names)):
        current_name = int(curr_field_names[i][3:])
        if one_hot_encoded_field == "MONTH":
            one_hot_df.rename(
                columns={curr_field_names[i]: new_name+month_names[current_name-1]}, inplace=True)
        elif one_hot_encoded_field == "WEEK_NUM_IN_MONTH" or one_hot_encoded_field == "WEEK_IN_MONTH":
            one_hot_df.rename(
                columns={curr_field_names[i]: new_name+week_names[current_name-1]}, inplace=True)
        elif (one_hot_encoded_field == "ITEM") | (one_hot_encoded_field == "CATEGORY_LEVEL_1"):
            one_hot_df.rename(
                columns={curr_field_names[i]: new_name+str(current_name)}, inplace=True)
        else:
            return None
    return one_hot_df


# Ensure train, validation, and test sets maintain the same data shape, same number of ohe fields
def add_missing_ohe_fields(one_hot_df,
                           cat_lvl_2_names,
                           item_num_names,
                           cat_lvl_1_names,
                           month_names=np.array(["January", "February", "March", "April", "May", "June",
                                                "July", "August", "September", "October", "November", "December"]),
                           week_names=np.array(
                               ["First_Week", "Second_Week", "Third_Week", "Fourth_Week", "Fifth_Week"])
                           ):
    """Adds any missing one hot encoded fields to one hot encoded dataframe"""
    for i_f in cat_lvl_2_names:
        if "is_"+i_f not in one_hot_df.columns:
            one_hot_df["is_"+i_f] = 0
    for w in week_names:
        if "is_"+w not in one_hot_df.columns:
            one_hot_df["is_"+w] = 0
    for m in month_names:
        if "is_"+m not in one_hot_df.columns:
            one_hot_df["is_"+m] = 0
    for i_n in item_num_names:
        if "is_"+i_n not in one_hot_df.columns:
            one_hot_df["is_"+i_n] = 0
    for i_c in cat_lvl_1_names:
        if "is_"+i_c not in one_hot_df.columns:
            one_hot_df["is_"+i_c] = 0
    return


# Split Data into X Dataframe and y array for independent variables and dependent variables, respectively
def split_data_X_y(data_set, target_var):
    """Returns X dataset and y target vector"""
    y_data_set = data_set[target_var]
    if isinstance(data_set, pd.DataFrame):
        data_set = DatasetFactory.from_dataframe(
            data_set.drop([target_var], axis=1), type_discovery=False)
    else:
        data_set = data_set.drop([target_var], axis=1)
    return data_set, y_data_set


# Get Data Split Proportion
def get_dataset_prop(X_train_ds, X_validation_ds, X_test_ds):
    """Return calculated train_validation_test split proportions"""
    total = X_train_ds.df.shape[0] + \
        X_validation_ds.df.shape[0] + X_test_ds.df.shape[0]
    train_prop = X_train_ds.df.shape[0] / total
    valid_prop = X_validation_ds.df.shape[0] / total
    test_prop = X_test_ds.df.shape[0] / total
    print("Training Set Proportion:    ", np.round(
        train_prop, 4), " Size: ", X_train_ds.df.shape[0])
    print("Validation Set Proportion:  ", np.round(
        valid_prop, 4), " Size: ", X_validation_ds.df.shape[0])
    print("Testing Set Proportion:     ", np.round(
        test_prop, 4), " Size: ", X_test_ds.df.shape[0])
    return (train_prop, valid_prop, test_prop)


# Drop categorical columns (covered by one hot encoding)
def drop_columns(data_set, columns_ls):
    """Drop selected column/list of columns"""
    data_set = data_set.df.drop(columns_ls, axis=1, inplace=False)
    data_set = DatasetFactory.from_dataframe(
        data_set, type_discovery=False)
    data_set.columns[:].ravel()
    return data_set


# Training data using a XGBRegressor model
def train_XGBRegressor(partition, gpu_available, train, y_train):
    """Return fitted XGBRegressor model"""
    if gpu_available:
        model = xgb.XGBRegressor(tree_method="gpu_hist", gpu_id=0,
                                 n_estimators=550, max_depth=8)
    else:
        model = xgb.XGBRegressor(n_estimators=550, max_depth=8)
    fit_start_time = time.time()
    print("Started Fitting Customer Site ",
          partition["full"]["CUSTOMER_SITE"].unique()[0])
    model.fit(train, y_train)
    partition["XGBmodel"] = model
    print("Fitting Completed")
    partition["XGBmodel_runtime"] = time.time() - fit_start_time

    print("Training Runtime (H:M:S):  {}".format(
        str(timedelta(seconds=(np.round(partition["XGBmodel_runtime"]))))))

    return partition["XGBmodel"]


# Get Feature Importance from training XGBRegressor model
def xgb_feature_importance_df(partition, train_columns):
    """Return dataframe of model feature importance"""
    partition["xgb_feature_importance"] = pd.DataFrame(data=partition["XGBmodel"].feature_importances_,
                                                       index=train_columns,
                                                       columns=["Importance"])
    return partition["xgb_feature_importance"]


# Training Score
def calc_training_error(partition, model, train, y_train):
    """Return model training score"""
    partition["training_score"] = model.score(train, y_train)
    print("Training Score:             ", np.round(
        partition["training_score"], 4))
    return partition["training_score"]


# Validation MSE
def calc_validation_mse(model, partition, validation, y_validation):
    """Return model validation MSE"""
    y_validation_pred = model.predict(validation)
    y_validation_pred.reshape(-1, 1)
    partition["y_validation_mse"] = MSE(y_validation, y_validation_pred)
    print("Validation MSE Score:       ", np.round(
        partition["y_validation_mse"], 4))
    return partition["y_validation_mse"]


# Validation RMSE
def calc_validation_rmse(model, partition, validation, y_validation):
    """Return model validation RMSE"""
    y_validation_pred = model.predict(validation)
    y_validation_pred.reshape(-1, 1)
    partition["y_validation_rmse"] = np.sqrt(
        MSE(y_validation, y_validation_pred))
    print("Validation RMSE Score:      ", np.round(
        partition["y_validation_rmse"], 4))
    return partition["y_validation_rmse"]


# Consolidate train/validation/test data with a new field "y_pred" for all predicted values
def calc_all_pred(model, partition, train, validation, test, y_train, y_validation, y_test
                  ):
    """Return dataframe with all predictions on full dataset"""
    # Customer Site predictions for each data split
    y_train_pred = model.predict(train)
    y_validation_pred = model.predict(validation)
    y_validation_pred.reshape(-1, 1)
    y_test_pred = model.predict(test)

    train.df["y_pred"] = y_train_pred
    validation.df["y_pred"] = y_validation_pred
    test.df["y_pred"] = y_test_pred

    train = pd.concat([train.df, y_train], axis=1)
    validation = pd.concat(
        [validation.df, y_validation], axis=1)
    test = pd.concat([test.df, y_test], axis=1)

    # Reset indices
    train.reset_index(inplace=True, drop=True)
    validation.reset_index(inplace=True, drop=True)
    test.reset_index(inplace=True, drop=True)

    # Combine all data into one dataframe
    partition["full"] = pd.concat([train, validation, test], axis=0)
    partition["full"] = pd.concat(
        [partition["full"], partition["categorical_data"]], axis=1)
    partition["full"].reset_index(drop=True, inplace=True)
    partition["full"] = partition["full"].sort_values(
        by=["DATE", "ITEM"], ascending=[True, True])

    return partition["full"], train, validation, test


# Delete dataframes and variables if needed
def remove_dataframes(partition, df_names):
    """Remove extraneous dataframes"""
    for name in df_names:
        del partition[name]
    gc.collect()
    return


# Create a dataframe that consists of all necessary fields for calculating mape
def mape_df(partition, data_set_str, grouped_field, columns):
    """Return dataset MAPE scores"""
    total_col = columns
    total_col.append(grouped_field)
    df_mape = partition[data_set_str][columns]
    if isinstance(df_mape, ADSDataset):
        df_mape = df_mape.df

    df_mape = df_mape[df_mape["SHIPMENTS_HISTORY"] != 0]
    df_mape[grouped_field] = partition[data_set_str][grouped_field]
    df_mape["abs_perc_error"] = abs(
        (df_mape["SHIPMENTS_HISTORY"] - df_mape["y_pred"])/df_mape["SHIPMENTS_HISTORY"])

    partition[data_set_str+"_mape"] = df_mape
    return (partition[data_set_str+"_mape"])


# Aggregate by mean to calculate MAPE at the "grouped_fields" level
def grouped_mape(partition, mape_df_str, grouped_fields):
    """Return Grouped MAPE scores"""
    if mape_df_str in partition:
        partition[mape_df_str+"_grouped"] = partition[mape_df_str].groupby(
            grouped_fields).mean().reset_index(drop = True).dropna()
        return partition[mape_df_str+"_grouped"]
    else:
        return None


def train_test_split_transformations(d, i, item_num_names, cat_lvl_1_names, cat_lvl_2_names):
    """Complete data transformations after train-validation-test split and train model as requested"""

    # Split data into training, validation, and test datasets
    train = create_data_split(
        partition=d[i], data="full",  max_date="2016-08-16")
    validation = create_data_split(
        partition=d[i], data="full", min_date="2016-08-16", max_date="2017-08-16")
    test = create_data_split(
        partition=d[i], data="full", min_date="2017-08-16")

    # One hot encode categorical fields in each of the three datasets
    for dataset in [train, validation, test]:
        item_family_one_hot = one_hot_categorical(
            data=dataset, field_name="CATEGORY_LEVEL_2")
        month_num_one_hot = one_hot_categorical(
            data=dataset, field_name="MONTH")
        week_num_month_one_hot = one_hot_categorical(
            data=dataset, field_name="WEEK_NUM_IN_MONTH")
        item_num_one_hot = one_hot_categorical(data=dataset, field_name="ITEM")
        item_class_one_hot = one_hot_categorical(
            data=dataset, field_name="CATEGORY_LEVEL_1")

        rename_one_hot_encoded_fields(
            month_num_one_hot, month_num_one_hot.columns, "MONTH")
        rename_one_hot_encoded_fields(
            week_num_month_one_hot, week_num_month_one_hot.columns, "WEEK_NUM_IN_MONTH")
        rename_one_hot_encoded_fields(
            item_num_one_hot, item_num_one_hot.columns, "ITEM")
        rename_one_hot_encoded_fields(
            item_class_one_hot, item_class_one_hot.columns, "CATEGORY_LEVEL_1")

        fields_ohe = pd.concat([item_family_one_hot, month_num_one_hot,
                               week_num_month_one_hot, item_num_one_hot, item_class_one_hot], axis=1)
        for x in fields_ohe.columns:
            dataset[x] = fields_ohe[x]

        add_missing_ohe_fields(one_hot_df=dataset,
                               month_names=np.array(["January", "February", "March", "April", "May", "June",
                                                    "July", "August", "September", "October", "November", "December"]),
                               week_names=np.array(
                                   ["First_Week", "Second_Week", "Third_Week", "Fourth_Week", "Fifth_Week"]),
                               cat_lvl_2_names=list(map(str, cat_lvl_2_names)),
                               item_num_names=list(map(str, item_num_names)),
                               cat_lvl_1_names=list(map(str, cat_lvl_1_names)))
        dataset = DatasetFactory.from_dataframe(
            dataset, target="SHIPMENTS_HISTORY", type_discovery=False)

    # Save original date and categorical fields for later use
    train_cat = train[["DATE", "ITEM", "CATEGORY_LEVEL_1",
                       "CATEGORY_LEVEL_2", "CUSTOMER_SITE", "CUSTOMER", "ZONE"]]
    validation_cat = validation[["DATE", "ITEM", "CATEGORY_LEVEL_1",
                                 "CATEGORY_LEVEL_2", "CUSTOMER_SITE", "CUSTOMER", "ZONE"]]
    test_cat = test[["DATE", "ITEM", "CATEGORY_LEVEL_1",
                     "CUSTOMER_SITE", "CATEGORY_LEVEL_2", "CUSTOMER", "ZONE"]]
    d[i]["categorical_data"] = pd.concat(
        [train_cat, validation_cat, test_cat], axis=0)

    # Split data into independent variables and predictor variable
    train, y_train = split_data_X_y(
        data_set=train, target_var="SHIPMENTS_HISTORY")
    validation, y_validation = split_data_X_y(
        data_set=validation, target_var="SHIPMENTS_HISTORY")
    test, y_test = split_data_X_y(
        data_set=test, target_var="SHIPMENTS_HISTORY")

    get_dataset_prop(X_train_ds=train,
                     X_validation_ds=validation, X_test_ds=test)

    # Drop categorical fields, leaving only numerical type fields
    train = drop_columns(data_set=train, columns_ls=["CATEGORY_LEVEL_2", "MONTH", "WEEK_NUM_IN_MONTH", "DATE",
                         "ITEM", "CATEGORY_LEVEL_1", "CUSTOMER_SITE", "CUSTOMER", "ZONE", "SALES_MONTH_YEAR", "SALES_WEEK_YEAR"])
    validation = drop_columns(data_set=validation, columns_ls=["CATEGORY_LEVEL_2", "MONTH", "WEEK_NUM_IN_MONTH", "DATE",
                              "ITEM", "CATEGORY_LEVEL_1", "CUSTOMER_SITE", "CUSTOMER", "ZONE", "SALES_MONTH_YEAR", "SALES_WEEK_YEAR"])
    test = drop_columns(data_set=test, columns_ls=["CATEGORY_LEVEL_2", "MONTH", "WEEK_NUM_IN_MONTH",  "DATE",
                        "ITEM", "CATEGORY_LEVEL_1", "CUSTOMER_SITE", "CUSTOMER", "ZONE", "SALES_MONTH_YEAR", "SALES_WEEK_YEAR"])

    # Store datasets for external use
    d[i]["training_col"] = (train.columns)
    d[i]["train"] = train
    d[i]["y_train"] = y_train
    d[i]["validation"] = validation
    d[i]["y_validation"] = y_validation
    test = test[d[i]["training_col"]]
    d[i]["test"] = test
    d[i]["y_test"] = y_test

    return


def train_model(d, i, fit_model=True, loaded_xgbmodel=None, gpu_available=False):
    if fit_model == False:
        # If not training model, returning the training and validation dataframes
        return d[i]["train"].df, d[i]["y_train"], d[i]["validation"].df, d[i]["y_validation"]
    elif fit_model and loaded_xgbmodel == None:
        # Train model with XGBRegressor and evaluate model performance
        train_XGBRegressor(
            partition=d[i], gpu_available=gpu_available, train=d[i]["train"], y_train=d[i]["y_train"])
    elif loaded_xgbmodel != None:
        # Use an existing model and evaluate model performance
        d[i]["XGBmodel"] = loaded_xgbmodel

    # Extract model's feature importance details
    xgb_feature_importance_df(
        partition=d[i], train_columns=d[i]["train"].columns)

    # Calculate evaluation metrics on training and validation datasets
    calc_training_error(
        partition=d[i], model=d[i]["XGBmodel"], train=d[i]["train"], y_train=d[i]["y_train"])
    calc_validation_mse(model=d[i]["XGBmodel"], partition=d[i],
                        validation=d[i]["validation"], y_validation=d[i]["y_validation"])
    calc_validation_rmse(model=d[i]["XGBmodel"], partition=d[i],
                         validation=d[i]["validation"], y_validation=d[i]["y_validation"])

    # Format full dataset (train, validation, test combined) to include all model predictions
    d[i]["full"], train, validation, test = calc_all_pred(model=d[i]["XGBmodel"], partition=d[i], train=d[i]["train"], validation=d[i]["validation"], test=d[i]["test"],
                                                          y_train=d[i]["y_train"], y_validation=d[i]["y_validation"], y_test=d[i]["y_test"])
    d[i]["full"] = d[i]["full"][["DATE", "ITEM", "CATEGORY_LEVEL_1", "CATEGORY_LEVEL_2", "CUSTOMER_SITE", "CUSTOM_MEASURE_1", 'WEEKEND', 'YEAR',
                                 'WEEK_NUM_IN_YEAR', 'DAY_NUM_IN_WEEK', 'HOLIDAY_INDICATOR_1', 'HOLIDAY_INDICATOR_2', 'LIFECYCLE_PHASE', 'EVENT_TYPE', "SHIPMENTS_HISTORY", "y_pred"]]

    return


def bayes_search_cv(partition, customer_site_ls, train, y_train, tss_n_splits, parameters, gpu_available, cv_n_iter):

    for i in customer_site_ls:
        if gpu_available:
            model = xgb.XGBRegressor(tree_method="gpu_hist", gpu_id=0)
        else:
            model = xgb.XGBRegressor()
        tscv = TimeSeriesSplit(n_splits=tss_n_splits)

        fit_start_time = time.time()
        print("Started Fitting: ")
        bayes_search = BayesSearchCV(model, parameters, cv=tscv, verbose=1, scoring=make_scorer(
            MSE, squared=False, greater_is_better=False), n_iter=cv_n_iter)
        bayes_result = bayes_search.fit(train.df, y_train)

        partition[i]["Bayes_best_score"] = bayes_result.best_score_
        partition[i]["Bayes_best_params"] = bayes_result.best_params_

        # Record runtime, best score, and best parameters
        bayes_timediff = time.time() - fit_start_time
        partition[i]["Bayes_runtime"] = bayes_timediff
        print("Training Runtime (H:M:S):  {}".format(
            str(timedelta(seconds=(np.round(bayes_timediff))))))

        print("Bayes Best Score: ", bayes_result.best_score_*-1)
        print("Bayes Best Params: ", bayes_result.best_params_)

    return bayes_result, bayes_timediff
