"""Some pre-processings steps implemented to prepare the bakery dataset for meta-learning"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/90_datasets/meta_bakery.ipynb.

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

# %% ../../nbs/90_datasets/meta_bakery.ipynb 3
import logging
logging.basicConfig(level=logging.INFO)

import numpy as np
import os
import pandas as pd

# %% ../../nbs/90_datasets/meta_bakery.ipynb 4
class BakeryDatasetLoader():

    """ Class to download the Kaggle M5 dataset and apply some preprocessing steps
    to prepare it for application in inventory management. """

    def __init__(self, data_path, overwrite=False, product_as_feature=False, store_as_features=False):
        self.create_paths(data_path)
        self.store_as_features = store_as_features
        self.product_as_feature = product_as_feature
    
    def load_dataset(self):

        """ Main function to load the dataset. """

        logging.info("Importing data")
        self.import_from_folder()
        logging.info("Preprocessing data")
        self.preprocess_pipeline()

        return self.demand, self.SKU_features, self.time_features, self.time_SKU_features, self.mask

    def create_paths(self, data_path):

        """ Create the paths for the data files. """

        self.data_path = data_path
        self.features_path = os.path.join(data_path, "bakery_data.csv")
        self.demand_path = os.path.join(data_path, "bakery_target.csv")

    def preprocess_pipeline(self):
    
        """ Apply simple preprocessing steps to the data. """
        
        logging.info("--Creating catogory mapping and features")  
        # self.sale["id"] = self.sale["id"].str.replace("_evaluation", "")
        unique_mapping = self.features[['store', 'product']].drop_duplicates()
        unique_mapping["store"] = unique_mapping["store"].astype(str)
        unique_mapping["product"] = unique_mapping["product"].astype(str)
        unique_mapping["SKU_id"] = unique_mapping["product"] + "_" + unique_mapping["store"]
        unique_mapping.set_index("SKU_id", inplace=True)

        dummy_columns = []
        if self.product_as_feature:
            dummy_columns.append("product")
        if self.store_as_features:
            dummy_columns.append("store")
        if dummy_columns:
            categories = pd.get_dummies(unique_mapping[dummy_columns], drop_first=True)
        else:
            categories = None
        
        logging.info("--Preparing calendric information")

        
        calendar_data = self.features[["date", "weekday", "month", "year"]]
        date_as_time = pd.to_datetime(calendar_data['date']).copy().unique()

        # get only rows with unique dates
        calendar_data = calendar_data.drop_duplicates()

        calendar_data['date'] = date_as_time

        # Adding the 'trend' column which counts from 0 and increments by 1
        calendar_data['trend'] = range(len(calendar_data))

        # Adding the 'quarter' column
        calendar_data['quarter'] = calendar_data['date'].dt.quarter

        # Defining the seasonal cut-off dates
        def get_season(date):
            year = date.year
            # Defining specific dates for season changes
            winter_start = pd.Timestamp(year=year, month=12, day=21)
            spring_start = pd.Timestamp(year=year, month=3, day=20)
            summer_start = pd.Timestamp(year=year, month=6, day=21)
            autumn_start = pd.Timestamp(year=year, month=9, day=22)
            
            if date >= winter_start or date < spring_start:
                return "Winter"
            elif spring_start <= date < summer_start:
                return "Spring"
            elif summer_start <= date < autumn_start:
                return "Summer"
            else:
                return "Autumn"

        # Adding the 'season' column with updated season logic
        calendar_data['season'] = calendar_data['date'].apply(get_season)

        dummy_columns = ["weekday", "month", "quarter", "season"]
        self.calendar = pd.get_dummies(calendar_data, columns=dummy_columns, drop_first=True)

        # drop date column
        self.calendar.drop(["date"], axis=1, inplace=True)

        # print(self.calendar)
        # for column in self.calendar.columns:
        #     print(column)

        logging.info("--Preparing demand")

        self.demand["store"] = self.features["store"].astype(str)   
        self.demand["product"] = self.features["product"].astype(str)   
        self.demand["date"] = self.features["date"]

        self.demand["SKU_index"] = self.demand["product"] + "_" + self.demand["store"]
        self.demand = self.demand.drop(["store", "product"], axis=1)

        self.demand = self.demand.pivot(index="date", columns="SKU_index", values="demand")

        # sort the columns
        self.demand = self.demand.reindex(sorted(self.demand.columns), axis=1)

        logging.info("--Preparing SKU-specific features")

        logging.info("--Preparing SKU-time-specific features")

        self.time_SKU_features = self.features.drop(["weekday", "month", "year", "is_holiday_next2days"], axis=1).copy()

        self.time_SKU_features["SKU_index"] = self.time_SKU_features["product"].astype(str) + "_" + self.time_SKU_features["store"].astype(str)
        self.time_SKU_features = self.time_SKU_features.drop(["store", "product"], axis=1)

        self.time_SKU_features.set_index(['date', 'SKU_index'], inplace=True)

        df_unstacked = self.time_SKU_features.unstack('SKU_index')

        df_unstacked = df_unstacked.sort_index(axis=0).sort_index(axis=1)

        self.time_SKU_features = df_unstacked

        self.time_SKU_features.sort_index(axis=1, level='SKU_index')
      
        self.SKU_features = categories # features that are not time-dependent
        self.time_features = self.calendar # features that are time-dependent
        self.time_SKU_features = self.time_SKU_features # features that are time- and SKU-dependent
        self.mask = None # A mask that can either mask datapoints during training or be used as a feature
         
    def import_from_folder(self):
        
        """ Import data from a folder. """

        self.features = pd.read_csv(self.features_path)
        self.demand = pd.read_csv(self.demand_path)
