# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_yaz.ipynb.

# %% ../nbs/01_yaz.ipynb 3
from __future__ import annotations
from fastcore.docments import *
from fastcore.test import *
from fastcore.utils import *

import pandas as pd
import numpy as np

from os.path import join
from tsfresh.feature_extraction import MinimalFCParameters
from tsfresh.utilities.dataframe_functions import roll_time_series
from tsfresh import extract_features
import pathlib
import pkg_resources

from .utils import *

# %% auto 0
__all__ = ['loadDataYaz']

# %% ../nbs/01_yaz.ipynb 5
def loadDataYaz(testDays = 28, returnXY = True, daysToCut = 0, unstacked = False):
    
    # LOAD DATA
    dataPath = pkg_resources.resource_stream(__name__, 'datasets/dataYaz_unprocessed.csv')
    data = pd.read_csv(dataPath)
    
    #---
    
    # DAY INDEX
    data = data.reset_index().rename(columns = {'index': 'dayIndex'})
    data['dayIndex'] = data.dayIndex

    #---
    
    # SEPARATE DEMAND OF DIFFERENT ITEMS
    X = data.iloc[:, 0:11]
    y = data.iloc[:, 12:]

    dataList = list()
    for col in y.columns:
        X_temp = X
        X_temp["item"] = col
        y_temp = y[col]
        data_temp = pd.concat([X_temp, y_temp], axis=1)
        data_temp.rename(columns={col: "demand"}, inplace=True)
        dataList.append(data_temp)

    data = pd.concat(dataList, axis = 0)

    #---
    
    # ID FEATURE AND SORTING
    data['id'] = data['item']
    data = data.sort_values(by = ['id', 'dayIndex'], axis = 0).reset_index(drop = True)

    #---

    # CUT DAYS DEPENDING ON DAYSTOCUT
    cutOffDate = data.dayIndex.max() - daysToCut
    data = data[data['dayIndex'] <= cutOffDate].reset_index(drop = True)
    
    #---
    
    # LABEL
    if isinstance(testDays, int):
        nDaysTest = testDays
    else:
        tsSizes = data.groupby(['id']).size()
        nDaysTest = int(tsSizes.iloc[0] * testDays)

    cutoffDateTest = data.dayIndex.max() - nDaysTest
    data['label'] = ['train' if data.dayIndex.iloc[i] <= cutoffDateTest else 'test' for i in range(data.shape[0])]    
    
    #---
    
    # NORMALIZE DEMAND
    scalingData = data[data.label == 'train'].groupby('id')['demand'].agg('max').reset_index()
    scalingData.rename(columns = {'demand': 'scalingValue'}, inplace = True)
    data = pd.merge(data, scalingData, on = 'id')

    data['demand'] = data.demand / data.scalingValue

    #---

    # DEMAND LAG FEATURES
    y = pd.DataFrame(data['demand'])
    X = data.drop(columns = ['demand'])

    # set lag features
    fc_parameters = MinimalFCParameters()

    # delete length features
    del fc_parameters['length']

    # create lag features
    X, y = add_lag_features(X = X, 
                            y = y, 
                            column_id = ['id'], 
                            column_sort = 'dayIndex', 
                            feature_dict = fc_parameters, 
                            time_windows = [(7, 7), (14, 14), (28, 28)],
                            n_jobs = 6, 
                            disable_progressbar = False)
    
    #---

    # CREATE UNSTACKED MULTIDIMENSIONAL DEMAND VECTOR IF DESIRED
    if unstacked:
        colsDemand = [column for column in X.columns if 'demand__' in column]
        colsOther = [column for column in X.columns if not 'demand__' in column]
        
        generalData = X[colsOther][X['id'] == X['id'][0]].reset_index(drop = True).drop('scalingValue', axis = 1)
        
        XList = list()
        yList = list()
        scalingValueList = list()
        
        XList.append(generalData)

        for item in np.unique(X['id']):
            XItem = X[X['id'] == item].reset_index(drop = True)
            yItem = y[X['id'] == item].reset_index(drop = True)

            newColNames = {col: col + '_' + item for col in colsDemand}
            
            XToAdd = XItem[colsDemand].rename(columns = newColNames)
            yToAdd = yItem.rename(columns = {'demand': 'demand_' + item})
            scalingValueToAdd = XItem[['scalingValue']].rename(columns = {'scalingValue': 'scalingValue_' + item})
            
            XList.append(XToAdd)
            yList.append(yToAdd)
            scalingValueList.append(scalingValueToAdd)
            
        X = pd.concat(XList, axis = 1)
        y = pd.concat(yList, axis = 1)
        scalingValues = pd.concat(scalingValueList, axis = 1)
        
        X = pd.concat([X, scalingValues], axis = 1)

    #---
    
    # DATE DUMMY VARIABLES
    X['year'] = X['year'].apply(lambda x: str(int(x)))

    X = pd.concat([X, 
                  pd.get_dummies(X.weekday, prefix = 'weekday'), 
                  pd.get_dummies(X.month, prefix = 'month'), 
                  pd.get_dummies(X.year, prefix = 'year')], axis = 1).drop(['weekday', 'month', 'year'], axis = 1)

    X = pd.concat([X, pd.get_dummies(X.item, prefix = 'item')], axis = 1).drop(['item'], axis = 1)

    #---
    
    # SPLIT INTO TRAIN AND TEST DATA
    data = pd.concat([y, X], axis = 1)
    XArray = np.array(X.drop(['label', 'id'], axis = 1))   
    
    if unstacked:
        yArray = np.array(y)           
    else:
        yArray = np.ravel(y)    
    
    XTrain = XArray[data['label'] == 'train']
    yTrain = yArray[data['label'] == 'train']

    XTest = XArray[data['label'] == 'test']
    yTest = yArray[data['label'] == 'test']

    #---

    if returnXY:
        return data, XTrain, yTrain, XTest, yTest
    else:
        return data    

