"""Dataloaders for tabular data"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/10_dataloaders/12_tabular_dataloaders.ipynb.

# %% auto 0
__all__ = ['XYDataLoader', 'MultiShapeLoader']

# %% ../../nbs/10_dataloaders/12_tabular_dataloaders.ipynb 3
import logging
logging.basicConfig(level=logging.INFO)

import numpy as np
from abc import ABC, abstractmethod
from typing import Union, Tuple, List, Literal
import pandas as pd
import math

from .base import BaseDataLoader

from sklearn.preprocessing import StandardScaler, MinMaxScaler

# %% ../../nbs/10_dataloaders/12_tabular_dataloaders.ipynb 4
class XYDataLoader(BaseDataLoader):

    """

    A class for datasets with the typicall X, Y structure. Both X
    and Y are numpy arrays. X may be of shape (datapoints, features) or (datapoints, sequence_length, features) 
    if lag features are used. The prep_lag_features can be used to create those lag features. Y is of shape
    (datapoints, units).

    """
    
    def __init__(self,
        X: np.ndarray,
        Y: np.ndarray,
        val_index_start: Union[int, None] = None, 
        test_index_start: Union[int, None] = None, 
        lag_window_params: Union[dict] = None, # default: {'lag_window': 0, 'include_y': False, 'pre_calc': False}
        normalize_features: Union[dict] = None, # default: {'normalize': True, 'ignore_one_hot': True}
    ):

        self.X = X
        self.Y = Y

        self.val_index_start = val_index_start
        self.test_index_start = test_index_start

        # train index ends either at the start of the validation set, the start of the test set or at the end of the dataset
        if self.val_index_start is not None:
            self.train_index_end = self.val_index_start-1
        elif self.test_index_start is not None:
            self.train_index_end = self.test_index_start-1
        else:
            self.train_index_end = len(Y)-1

        self.dataset_type = "train"

        normalize_features = normalize_features or {'normalize': True, 'ignore_one_hot': True}
        lag_window_params = lag_window_params or {'lag_window': 0, 'include_y': False, 'pre_calc': False}

        self.normalize_features(**normalize_features, initial_normalization=True)
        self.prep_lag_features(**lag_window_params)

        # X must at least have datapoint and feature dimension
        if len(X.shape) == 1:
            self.X = X.reshape(-1, 1)
        
        # Y must have at least datapoint and unit dimension (even if only one unit is present)
        if len(Y.shape) == 1:
            self.Y = Y.reshape(-1, 1)

        assert len(X) == len(Y), 'X and Y must have the same length'

        self.num_units = Y.shape[1] # shape 0 is alsways time, shape 1 is the number of units (e.g., SKUs)

        super().__init__()

    def normalize_features(self,
        normalize: bool = True,
        ignore_one_hot: bool = True,
        initial_normalization=False # Flag if it is set before having added lag features
        ):

        """
        Normalize features using a standard scaler. If ignore_one_hot is true, one-hot encoded features are not normalized.

        """

        if normalize:

            scaler = StandardScaler()

            if initial_normalization:

                if len(self.X.shape) == 3:
                    raise ValueError('Normalization not possible with lag features. Please set initial_normalization=False')
            
                scaler.fit(self.X[:self.train_index_end+1]) # +1 to include the last training point
                scaler.transform(self.X)

                if initial_normalization:
                    return
                else:
                    raise NotImplementedError('Normalization after lag features have been set not implemented yet')

                    # Idea:
                        # remove time dimension
                        # normalize features
                        # add time_dimension back
                    # Problem:
                        # usage of prep_lag_features needs to ensure y is not added a second time

    def prep_lag_features(self,
        lag_window: int = 0, # length of the lage window
        include_y: bool = False, # if lag demand shall be included as feature
        pre_calc: bool = False # if all lags are pre-calculated for the entire dataset
        ):

        """
        Create lag feature for the dataset. If "inlcude_y" is true, then a lag-1 of of the target variable is added as a feature.
        If lag-window is > 0, the lag features are added as middle dimension to X. Note that this, e.g., means that with a lag
        window of 1, the data will include 2 time steps, the current features including lag-1 demand and the lag-1 features
        including lag-2 demand. If pre-calc is true, all these calculations are performed on the entire dataset reduce
        computation time later on at the expense of increases memory usage. 

        """
        # to be discussed: Do we need option to only provide lag demand wihtout lag features?
        self.lag_window = lag_window
        self.pre_calc = pre_calc
        self.include_y = include_y
        
        if self.pre_calc:
            if self.include_y:
                # add additional column to X with demand shifted by 1
                self.X = np.concatenate((self.X, np.roll(self.Y, 1, axis=0)), axis=1)
                self.X = self.X[1:] # remove first row
                self.Y = self.Y[1:] # remove first row
                
                self.val_index_start = self.val_index_start-1
                self.test_index_start = self.test_index_start-1
                self.train_index_end  = self.train_index_end-1
        
            if self.lag_window is not None and self.lag_window > 0:

                # add lag features as dimention 2 to X (making it dimension (datapoints, sequence_length, features))
                X_lag = np.zeros((self.X.shape[0], self.lag_window+1, self.X.shape[1]))
                for i in range(self.lag_window+1):
                    if i == 0:
                        features = self.X
                    else:    
                        features = self.X[:-i, :]
                    X_lag[i:, self.lag_window-i, :] = features
                self.X = X_lag[self.lag_window:]
                self.Y = self.Y[self.lag_window:]

                self.val_index_start = self.val_index_start-self.lag_window
                self.test_index_start = self.test_index_start-self.lag_window
                self.train_index_end  = self.train_index_end-self.lag_window

        else:
            self.lag_window = None
            self.include_y = False
            # add time dimension to X

    def update_lag_features(self,
        lag_window: int,
        ):

        """ Update lag window parameters for dataloader object that is already initialized """

        raise NotImplementedError('Not implemented yet')

        # Problem: updating lag_features naively would shorten the dataset each time it is called

    def __getitem__(self, idx): 

        """ get item by index, depending on the dataset type (train, val, test)"""

        if self.dataset_type == "train":
            if idx > self.train_index_end:
                raise IndexError(f'index {idx} out of range{self.train_index_end}')
            idx = idx

        elif self.dataset_type == "val":
            idx = idx + self.val_index_start
            
            if idx >= self.test_index_start:
                raise IndexError(f'index{idx} out of range{self.test_index_start}')
            
        elif self.dataset_type == "test":
            idx = idx + self.test_index_start
            
            if idx >= len(self.X):
                raise IndexError(f'index{idx} out of range{len(self.X)}')
        
        else:
            raise ValueError('dataset_type not set')

        return self.X[idx], self.Y[idx]

    def __len__(self):
        return len(self.X)
    
    @property
    def X_shape(self):
        return self.X.shape
    
    @property
    def Y_shape(self):
        return self.Y.shape

    @property
    def len_train(self):
        return self.train_index_end+1

    @property
    def len_val(self):
        if self.val_index_start is None:
            raise ValueError('no validation set defined')
        return self.test_index_start-self.val_index_start

    @property
    def len_test(self):
        if self.test_index_start is None:
            raise ValueError('no test set defined')
        return len(self.Y)-self.test_index_start

    def get_all_X(self,
                dataset_type: str = 'train' # can be 'train', 'val', 'test', 'all'
                ): 

        """
        Returns the entire features dataset.
        Return either the train, val, test, or all data.
        """

        if dataset_type == 'train':
            return self.X[:self.val_index_start].copy() if self.X is not None else None
        elif dataset_type == 'val':
            return self.X[self.val_index_start:self.test_index_start].copy() if self.X is not None else None
        elif dataset_type == 'test':
            return self.X[self.test_index_start:].copy() if self.X is not None else None
        elif dataset_type == 'all':
            return self.X.copy() if self.X is not None else None
        else:
            raise ValueError('dataset_type not recognized')

    def get_all_Y(self,
                dataset_type: str = 'train' # can be 'train', 'val', 'test', 'all'
                ): 

        """
        Returns the entire target dataset.
        Return either the train, val, test, or all data.
        """

        if dataset_type == 'train':
            return self.Y[:self.val_index_start].copy() if self.Y is not None else None
        elif dataset_type == 'val':
            return self.Y[self.val_index_start:self.test_index_start].copy() if self.Y is not None else None
        elif dataset_type == 'test':
            return self.Y[self.test_index_start:].copy() if self.Y is not None else None
        elif dataset_type == 'all':
            return self.Y.copy() if self.Y is not None else None
        else:
            raise ValueError('dataset_type not recognized')
        

# %% ../../nbs/10_dataloaders/12_tabular_dataloaders.ipynb 18
class MultiShapeLoader(BaseDataLoader):

    """
    A class designed for comlex datasets with mutlipe feature types. The class is more
    memory-efficient than the XYDataLoader, as it separate the storeage of SKU-specific
    feature, time-specific features, and time-SKU-specific features. The class works generically
    as long as those feature classes are provided during pre-processing. The class is designed 
    to handle classic learning, but able to work in a meta-learning pipeline where no SKU-dimension
    is present and the model needs to make prediction on SKU-time level without knowhing the
    specific SKU.
    """
    
    def __init__(self,
        # mandatory data
        demand: pd.DataFrame, # Demand data of shape time x SKU
        time_features: pd.DataFrame, # Features constant over SKU of shape time x time_features
        time_SKU_features: pd.DataFrame, # Features varying over time and SKU of shape time x (time_SKU_features*SKU) with double index
        
        # optional data
        mask: pd.DataFrame = None, # Mask of shape time x SKU telling which SKUs are available at which time (can be used as mask during trainig or added to features)
        SKU_features: pd.DataFrame = None, # Features constant over time of shape SKU x SKU_features - only for algorithms learning across SKUs
        
        val_index_start: Union[int, None] = None, # Validation index start on the time dimension
        test_index_start: Union[int, None] = None, # Test index start on the time dimension
        in_sample_val_test_SKUs: List = None, # SKUs in the training set to be used for validation and testing, out-of-sample w.r.t. time dimension
        out_of_sample_val_SKUs: List = None, # SKUs to be hold-out for validation (can be same as test if no validation on out-of-sample SKUs required)
        out_of_sample_test_SKUs: List = None, # SKUs to be hold-out for testing
        lag_window_params: dict | None = None, # default: {'lag_window': 0, 'include_y': False, 'pre_calc': True}
        normalize_features: dict | None = None, # default: {'normalize': True, 'ignore_one_hot': True}
        engineered_SKU_features: Union[dict] = None, # default: ["mean_demand", "std_demand", "kurtosis_demand", "skewness_demand", "percentile_10_demand", "percentile_30_demand", "median_demand", "percentile_70_demand", "percentile_90_demand", "inter_quartile_range"]
        use_engineered_SKU_features: bool = False, # if engineered features shall be used
        include_non_available: bool = False, # if timestep/SKU combination where the SKU was not available for sale shall be included. If included, it will be used as feature, otherwise as mask.
        train_subset: int = None ,# if only a subset of SKUs is used for training. Will always contain in_sample_val_test_SKUs and then fills the rest with random SKUs
        train_subset_SKUs: List = None, # if train_subset is set, specific SKUs can be provided
        meta_learn_units: bool = False, # if units (SKUs) are trained in the batch dimension to meta-learn across SKUs
        lag_demand_normalization: Literal['minmax', 'standard', 'no_normalization'] | None = "standard", # minmax, standard, no_normalization or None. If None, same demand_normalization
        demand_normalization: Literal['minmax', 'standard', 'no_normalization'] = 'no_normalization', # 'standard' or 'minmax'
        demand_unit_size: float | None = None, # use same convention as for other dataloaders and enviornments, but here only full decimal values are allowed
        provide_additional_target: bool = False, # follows ICL convention by providing actual demand to token, with the last token receiving 0
        permutate_inputs: bool = False, # if the inputs shall be permutated during training for meta-learning
    ):
     
        logging.info("Setting main env attributes")
        # Set data
        self.demand = demand
        self.SKU_features = SKU_features
        self.SKU_features_out_of_sample_val, self.SKU_features_out_of_sample_test = None, None # will be overwritten if relevant
        self.time_features = time_features
        self.time_SKU_features = time_SKU_features
        self.mask = mask
        self.permutate_inputs = permutate_inputs

        # convert dtypes to float
        self.demand = self.demand.astype(float)
        self.time_features = self.time_features.astype(float)
        self.time_SKU_features = self.time_SKU_features.astype(float)
        if self.SKU_features is not None:
            self.SKU_features = self.SKU_features.astype(float)
        if self.mask is not None:
            self.mask = self.mask.astype(float)

        # Set default values for dict inputs:
        normalize_features = normalize_features or {'normalize': True, 'ignore_one_hot': True}
        lag_window_params = lag_window_params or {'lag_window': 0, 'include_y': False, 'pre_calc': False}
        self.lag_window_params = lag_window_params # lag window parameters saved as attribute
        engineered_SKU_features = engineered_SKU_features or ["mean_demand", "std_demand", "kurtosis_demand", "skewness_demand", "percentile_10_demand", "percentile_30_demand", "median_demand", "percentile_70_demand", "percentile_90_demand", "inter_quartile_range"]
        if not use_engineered_SKU_features:
            engineered_SKU_features = None

        # Set further training parameters
        self.include_non_available = include_non_available
        self.meta_learn_units = meta_learn_units
        train_subset, train_subset_SKUs = self.set_train_subset(train_subset, train_subset_SKUs) # set the attributes train_subset and train_subset_SKUs
        self.lag_demand_normalization = lag_demand_normalization if lag_demand_normalization is not None else demand_normalization
        self.demand_normalization = demand_normalization
        if demand_unit_size is not None:
            self.demand_unit_size = int(-math.log10(demand_unit_size))
            if self.demand_unit_size % 1 != 0:
                raise ValueError('demand_unit_size must be a full decimal value')
            else:
                self.demand_unit_size = int(self.demand_unit_size)
        else:
            self.demand_unit_size = None
        self.provide_additional_target = provide_additional_target
        
        # Some necessary flags
        ## Whether the features are already normalized
        self.normalized_in_sample_SKUs = False
        self.normalized_out_of_sample_SKUs = False
        self.added_engineereed_features_to_in_sample_SKUs = False
        self.added_engineereed_features_to_out_of_sample_val_SKUs = False
        self.added_engineereed_features_to_out_of_sample_test_SKUs = False
        ## What data to return in the __getitem__ method
        self.SKU_type = "in_sample" # or "out_of_sample_val" or "out_of_sample_test" # What kind of SKU retruning in __getitem__
        self.dataset_type = "train" # or "val" or "test", affecting the time-dimension

        logging.info("Setting indices for validation and test set")
        # Determine train, val, test indices (time dimension only!)
        self.val_index_start = val_index_start
        self.test_index_start = test_index_start
        self.train_index_start = self.lag_window_params["lag_window"] # start index for training data later to allow for lag features
        self.train_index_start += self.lag_window_params["include_y"] # if lag demand is included as feature need one more timestep

        # train index ends either at the start of the validation set, the start of the test set or at the end of the dataset
        if self.val_index_start is not None:
            self.train_index_end = self.val_index_start-1
        elif self.test_index_start is not None:
            self.train_index_end = self.test_index_start-1
        else:
            self.train_index_end = len(self.demand)-1

        logging.info("Setting out-of-sample SKUs")
        # set out-of-sample SKUs - note that this is not done by indecing, but the SKUs data will be
        # removed from the in-sample data and stored in separate attributes. Therefore, there are no
        # index attributes for out-of-sample SKUs. The feature attributes are stored by the function below.
        self.out_of_sample_val_SKUs, self.out_of_sample_test_SKUs = self.set_out_of_sample_SKUs(out_of_sample_val_SKUs, out_of_sample_test_SKUs)

        if self.out_of_sample_val_SKUs is not None and self.out_of_sample_test_SKUs is not None:
            if len(self.out_of_sample_val_SKUs) != 0 or len(self.out_of_sample_test_SKUs) != 0:
                self.out_of_sample = True
            else:
                self.out_of_sample = False
        else:
            self.out_of_sample = False

        logging.info("Identifying training and in-sample validation and test SKUs")
        self.in_sample_val_test_SKUs = self.set_in_sample_val_test_SKUs(in_sample_val_test_SKUs) # need to determine index later after out-of-sample units are removed
        self.train_SKUs = self.identify_train_SKUs(train_subset, train_subset_SKUs) # need to determine index later after out-of-sample units are removed 
        if self.in_sample_val_test_SKUs is None:
            self.in_sample_val_test_SKUs = self.train_SKUs # validation and training SKUs are teh same

        logging.info("Determine number of units and features")
        # Num units is relevant for the output dimension when validating and testing. If the model is not trained as a 
        # meta learner it is identicall in traing, and validation/testing. If it is trained as a meta learner, the output
        # dimension is the number of in_sample_val_test_SKUs, iregardless of the number of SKUs in the training set.
        self.num_units = len(self.in_sample_val_test_SKUs) if self.in_sample_val_test_SKUs is not None else len(self.demand.columns)
        
        # Determine number of features
        self.num_time_SKU_features = len(self.time_SKU_features.columns.get_level_values(0).unique())
        if lag_window_params["include_y"]:
            self.num_time_SKU_features += 1 # lag demand is time and SKU specific
        if include_non_available:
            self.num_time_SKU_features += 1 # non-availability is time and SKU specific
        if self.provide_additional_target:
            self.num_time_SKU_features += 1
        self.num_time_features = len(self.time_features.columns)
        self.num_SKU_features = len(self.SKU_features.columns) if self.SKU_features is not None else 0
        if engineered_SKU_features is not None:
            self.num_SKU_features += len(engineered_SKU_features) # engineered features are SKU specific

        # print("xxxxxxx Number features per type:")
        # print("num_time_SKU_features:", self.num_time_SKU_features)
        # print("num_time_features:", self.num_time_features)
        # print("num_SKU_features:", self.num_SKU_features)

        # Determine total number of features per datapoint
        self.num_features = self.num_SKU_features + self.num_time_features + self.num_time_SKU_features

        if engineered_SKU_features is not None:
            logging.info("Creating engineered SKU features for training data")
            engineered_SKU_features_data = self.build_engineered_SKU_features(engineered_SKU_features, self.demand.iloc[:self.train_index_end+1]) # only for training data initially
            self.SKU_features = pd.concat([self.SKU_features, engineered_SKU_features_data.transpose()], axis=1)
            self.added_engineereed_features_to_in_sample_SKUs = True

            logging.info("Creating engineered SKU features for out-of-sample validation data")
            engineered_SKU_features_data = self.build_engineered_SKU_features(engineered_SKU_features, self.demand_out_of_sample_val) # only for out-of-sample validation data
            self.SKU_features_out_of_sample_val = pd.concat([self.SKU_features_out_of_sample_val, engineered_SKU_features_data.transpose()], axis=1)
            self.added_engineereed_features_to_out_of_sample_val_SKUs = True

            logging.info("Creating engineered SKU features for out-of-sample test data")
            engineered_SKU_features_data = self.build_engineered_SKU_features(engineered_SKU_features, self.demand_out_of_sample_test) # only for out-of-sample test data
            self.SKU_features_out_of_sample_test = pd.concat([self.SKU_features_out_of_sample_test, engineered_SKU_features_data.transpose()], axis=1)
            self.added_engineereed_features_to_out_of_sample_test_SKUs = True

        logging.info("Normalizing in-sample SKU features (based on training timesteps)")
        self.normalize_demand_and_features_in_sample(normalize=normalize_features["normalize"], ignore_one_hot=normalize_features["ignore_one_hot"],initial_normalization=True)

        if self.out_of_sample:
            logging.info("Normalizing out-of-sample SKU features")
            self.normalize_demand_and_features_out_of_sample(normalize=normalize_features["normalize"], ignore_one_hot=normalize_features["ignore_one_hot"],initial_normalization=True)

        logging.info("Saving data as numpy and saving indices")
        # store row and column indices of demand, SKU_features time_features mask and then convert to numpy array
        logging.info("--Saving indices")

        ############ in sample data ############
        self.in_sample_val_test_SKUs_indices = self.demand.columns.get_indexer(self.in_sample_val_test_SKUs) if self.in_sample_val_test_SKUs is not None else None

        self.train_SKUs_indices = self.demand.columns.get_indexer(self.train_SKUs)

        self.demand_indices = self.save_indices(self.demand)
        self.SKU_features_indices = self.save_indices(self.SKU_features) if self.SKU_features is not None else None
        self.time_features_indices = self.save_indices(self.time_features)
        self.time_SKU_features_indices = self.save_indices(self.time_SKU_features)
        self.mask_indices = self.save_indices(self.mask)

        logging.info("--Converting to numpy - in sample")
        self.demand = self.demand.to_numpy()
        self.demand_lag = self.demand_lag.to_numpy()
        self.SKU_features = self.SKU_features.to_numpy() if self.SKU_features is not None else None
        self.time_features = self.time_features.to_numpy()
        self.time_SKU_features = self.time_SKU_features.to_numpy()
        self.mask = self.mask.to_numpy()

        ############ out of sample data ############

        if self.out_of_sample:
            self.out_of_sample_val_SKUs_indices = self.demand_out_of_sample_val.columns.get_indexer(self.out_of_sample_val_SKUs)
            self.out_of_sample_test_SKUs_indices = self.demand_out_of_sample_test.columns.get_indexer(self.out_of_sample_test_SKUs)

            self.demand_out_of_sample_val_indices = self.save_indices(self.demand_out_of_sample_val)
            self.SKU_features_out_of_sample_val_indices = self.save_indices(self.SKU_features_out_of_sample_val) if self.SKU_features_out_of_sample_val is not None else None
            self.time_SKU_features_out_of_sample_val_indices = self.save_indices(self.time_SKU_features_out_of_sample_val)
            self.mask_out_of_sample_val_indices = self.save_indices(self.mask_out_of_sample_val)
        
            logging.info("--Converting to numpy - out of sample val")
            self.demand_out_of_sample_val = self.demand_out_of_sample_val.to_numpy()
            self.demand_lag_out_of_sample_val = self.demand_lag_out_of_sample_val.to_numpy()
            self.SKU_features_out_of_sample_val = self.SKU_features_out_of_sample_val.to_numpy() if self.SKU_features_out_of_sample_val is not None else None
            self.time_SKU_features_out_of_sample_val = self.time_SKU_features_out_of_sample_val.to_numpy()
            self.mask_out_of_sample_val = self.mask_out_of_sample_val.to_numpy()

            self.demand_out_of_sample_test_indices = self.save_indices(self.demand_out_of_sample_test)
            self.SKU_features_out_of_sample_test_indices = self.save_indices(self.SKU_features_out_of_sample_test) if self.SKU_features_out_of_sample_test is not None else None
            self.time_SKU_features_out_of_sample_test_indices = self.save_indices(self.time_SKU_features_out_of_sample_test)
            self.mask_out_of_sample_test_indices = self.save_indices(self.mask_out_of_sample_test)

            logging.info("--Converting to numpy - out of sample test")
            self.demand_out_of_sample_test = self.demand_out_of_sample_test.to_numpy()
            self.demand_lag_out_of_sample_test = self.demand_lag_out_of_sample_test.to_numpy()
            self.SKU_features_out_of_sample_test = self.SKU_features_out_of_sample_test.to_numpy() if self.SKU_features_out_of_sample_test is not None else None
            self.time_SKU_features_out_of_sample_test = self.time_SKU_features_out_of_sample_test.to_numpy()
            self.mask_out_of_sample_test = self.mask_out_of_sample_test.to_numpy()

        

        ############ final params ############
        self.len_train_time = self.train_index_end-self.train_index_start+1

        print("len_train_time:", self.len_train_time)
        print("num_units:", self.num_units)
        print("len train_SKU_indices:", len(self.train_SKUs_indices))

        if self.meta_learn_units:
            logging.info("--Creating time-SKU index for training data")
            self.sku_time_index = [(i, j) for i in self.train_SKUs_indices for j in range(self.len_train_time)]

        self.set_return_sku("in_sample")

        super().__init__()

    def set_train_subset(self, train_subset, train_subset_SKUs):
        """ Prepare setting the attributes train_subset and train_subset_SKUs """

        if train_subset is not None and train_subset < 1:
            raise ValueError('train_subset must be an integer larger than 0')

        if train_subset_SKUs is not None and not isinstance(train_subset_SKUs, list):
            raise ValueError(f'train_subset_SKUs must be a list of SKUs, got {type(train_subset_SKUs)}')

        if train_subset_SKUs is not None:
            if train_subset is None:
                logging.info("--Infering train_subset from train_subset_SKUs")
                train_subset = len(train_subset_SKUs)
            else:
                if train_subset != len(train_subset_SKUs):
                    raise ValueError('train_subset_SKUs must have the same length as train_subset')
        
        return train_subset, train_subset_SKUs

    def set_in_sample_val_test_SKUs(self, in_sample_val_test_SKUs):

        if in_sample_val_test_SKUs is None:
            in_sample_val_test_SKUs_indices = None
        else:
            if not self.meta_learn_units:
                raise ValueError('in_sample_val_test_SKUs can only be set if meta_learn_units is True/n \
                otherwise the output dimension needs to be the same for training and validation/testing')
            if not set(in_sample_val_test_SKUs).issubset(self.demand.columns):

                missing_SKUs = 0
                names = []
                for i in in_sample_val_test_SKUs:
                    if i not in self.demand.columns:
                        missing_SKUs += 1
                        names.append(i)
                
                print(names)
            

                raise ValueError('in_sample_val_test_SKUs must be a subset of all SKUs')

        return in_sample_val_test_SKUs
            
    def identify_train_SKUs(self, train_subset, train_subset_SKUs):
        """ determine which SKUs are used for training, validation and testing """
        
        if train_subset is not None:

            if train_subset_SKUs is not None:
                # check that all train_SKUs are in demand.collumns
                if not set(train_subset_SKUs).issubset(self.demand.columns):
                    raise ValueError('train_subset_SKUs must be a subset of all training SKUs')
                if self.in_sample_val_test_SKUs is not None:
                    if not set(self.in_sample_val_test_SKUs).issubset(train_subset_SKUs):
                        raise ValueError('train_subset_SKUs must contain in_sample_val_test_SKUs')
            else:
                if self.in_sample_val_test_SKUs is not None and train_subset < len(self.in_sample_val_test_SKUs):
                    raise ValueError(f'train_subset ({train_subset}) must be equal or larger than the number of in_sample_val_test_SKUs ({len(self.in_sample_val_test_SKUs)})')
                train_subset_SKUs = self.in_sample_val_test_SKUs if self.in_sample_val_test_SKUs is not None else []
                remaining_SKUs = self.demand.columns.difference(train_subset_SKUs)
                additional_SKUs = np.random.choice(remaining_SKUs, train_subset-len(train_subset_SKUs), replace=False)
                train_SKUs = np.concatenate((train_subset_SKUs, additional_SKUs))
                # order train_SKUs as in demand.columns
                train_subset_SKUs = [sku for sku in self.demand.columns if sku in train_SKUs]
        else:
            train_subset_SKUs = self.demand.columns # val and test SKUs have been removed before, only training SKUs remain 

        return train_subset_SKUs
    
    def set_out_of_sample_SKUs(self,
            out_of_sample_val_SKUs: List | np.ndarray | None, # SKUs for validation
            out_of_sample_test_SKUs: List | np.ndarray | None, # SKUs for testing. If same dataset shall be used, it can be identical
            ):

        """
        Sets out-of-sample SKUs for validation and test datasets, and removes them from the in-sample data.
        It handles the cases that no out-of-sample SKUs are set (both inputs are None), one of them is set,
        both of them are set to different values, and both of them are identical. In-sample data remain,
        while out-of-sample data are moved to seperate attributes.

        """

        # Note: one of the two might be None, both might be None, or both might be identical

        self.test_out_of_sample_SKUs(out_of_sample_val_SKUs, out_of_sample_test_SKUs)

        if out_of_sample_val_SKUs is None:
            out_of_sample_val_SKUs = []
        if out_of_sample_test_SKUs is None:
            out_of_sample_test_SKUs = []

        for sku, attr_suffix in [(out_of_sample_val_SKUs, 'val'), (out_of_sample_test_SKUs, 'test')]:
            if sku is not None:
                logging.info(f"--Setting out-of-sample {attr_suffix} SKUs")
                # Set attributes for out-of-sample SKUs
                setattr(self, f'demand_out_of_sample_{attr_suffix}', self.demand.loc[:, sku])
                if self.SKU_features is not None:
                    setattr(self, f'SKU_features_out_of_sample_{attr_suffix}', self.SKU_features.loc[sku])
                setattr(self, f'time_SKU_features_out_of_sample_{attr_suffix}', # here SKU are in columns on index level 2
                        self.time_SKU_features.loc[:, pd.IndexSlice[:, sku]]) 
                setattr(self, f'mask_out_of_sample_{attr_suffix}', self.mask.loc[:, sku])
                # time_features are independent of SKU, so no need to set them

        # unique values of uniion of both lists
        skus_to_remove = list(set(out_of_sample_val_SKUs).union(out_of_sample_test_SKUs))
        if skus_to_remove:
            logging.info(f"--Removing {len(skus_to_remove)} out-of-sample SKUs from in-sample data")
            # Remove out-of-sample SKUs from in-sample data
            self.demand.drop(columns=skus_to_remove, inplace=True)
            if self.SKU_features is not None:
                self.SKU_features.drop(index=skus_to_remove, inplace=True)
            for single_sku in skus_to_remove:
                columns_to_drop = self.time_SKU_features.columns.get_loc_level(single_sku, level=1)
                self.time_SKU_features.drop(columns=self.time_SKU_features.columns[columns_to_drop[0]], inplace=True)
            self.mask.drop(columns=skus_to_remove, inplace=True)

        return out_of_sample_val_SKUs, out_of_sample_test_SKUs

    def test_out_of_sample_SKUs(self, out_of_sample_val_SKUs, out_of_sample_test_SKUs):
        
        """
        Validates the out-of-sample SKU lists for compatibility. Ensures that the inputs are either None, lists, or numpy arrays,
        and that they meet specific conditions regarding uniqueness and overlap.
        """

        # Check if inputs are either None, list, or numpy array
        if not (isinstance(out_of_sample_val_SKUs, (list, np.ndarray)) or out_of_sample_val_SKUs is None):
            raise ValueError("out_of_sample_val_SKUs must be a list, numpy array, or None.")
        if not (isinstance(out_of_sample_test_SKUs, (list, np.ndarray)) or out_of_sample_test_SKUs is None):
            raise ValueError("out_of_sample_test_SKUs must be a list, numpy array, or None.")

        # Convert None to empty list for easier processing
        if out_of_sample_val_SKUs is None:
            out_of_sample_val_SKUs = []
        if out_of_sample_test_SKUs is None:
            out_of_sample_test_SKUs = []

        # Ensure all entries are unique within each list
        if len(out_of_sample_val_SKUs) != len(set(out_of_sample_val_SKUs)):
            raise ValueError("out_of_sample_val_SKUs contains duplicate entries.")
        if len(out_of_sample_test_SKUs) != len(set(out_of_sample_test_SKUs)):
            raise ValueError("out_of_sample_test_SKUs contains duplicate entries.")

        # Ensure there are no overlapping SKUs between validation and test sets (unless they are meant to be identical)

        if out_of_sample_val_SKUs is not None and out_of_sample_test_SKUs is not None:
            overlap = set(out_of_sample_val_SKUs).intersection(set(out_of_sample_test_SKUs))

            if len(overlap) > 0:
                if len(overlap) == len(out_of_sample_val_SKUs):
                    logging.warning("--Out of sample validation and test SKUs are identical - ensure this is intended.")
                else:
                    raise ValueError(f"Validation and test SKUs must not overlap. Overlapping SKUs: {overlap}")

        logging.debug("--Out-of-sample SKUs have passed validation checks.")
        
    @staticmethod
    def build_engineered_SKU_features(engineered_SKU_features: List, demand: pd.DataFrame):

        """
        Create engineered features for each SKU
        """

        feature_names = []
        feature_values = []

        for feature in engineered_SKU_features:

            if feature == "mean_demand":
                mean_demand = demand.mean(axis=0)
            elif feature == "std_demand":
                std_demand = demand.std(axis=0)
            elif feature == "kurtosis_demand":
                kurtosis_demand = demand.kurtosis(axis=0)
            elif feature == "skewness_demand":
                skewness_demand = demand.skew(axis=0)
            elif feature == "percentile_10_demand":
                percentile_10_demand = demand.quantile(0.1, axis=0)
            elif feature == "percentile_30_demand":
                percentile_30_demand = demand.quantile(0.3, axis=0)
            elif feature == "median_demand":
                median_demand = demand.median(axis=0)
            elif feature == "percentile_70_demand":
                percentile_70_demand = demand.quantile(0.7, axis=0)
            elif feature == "percentile_90_demand":
                percentile_90_demand = demand.quantile(0.9, axis=0)
            elif feature == "inter_quartile_range":
                inter_quartile_range = demand.quantile(0.75, axis=0) - demand.quantile(0.25, axis=0)
            else:  
                raise ValueError(f'Feature {feature} not recognized')
            
            feature_names.append(feature)
            feature_values.append(locals()[feature])

        return pd.DataFrame(feature_values, columns=demand.columns, index=feature_names)
    
    def normalize_demand_and_features_in_sample(self,
        normalize: bool = True,
        ignore_one_hot: bool = True,
        initial_normalization = False # Flag if it is set before having added lag features
        ):

        """
        Normalize features using a standard scaler. If ignore_one_hot is true, one-hot encoded features are not normalized.
        """

        if normalize:

            if self.normalized_in_sample_SKUs:
                raise ValueError('Features already normalized')

            if self.demand_normalization == 'minmax':
                self.scaler_demand = MinMaxScaler()
                self.scaler_out_of_sample_val_demand = MinMaxScaler()
                self.scaler_out_of_sample_test_demand = MinMaxScaler()
            elif self.demand_normalization == 'standard':
                self.scaler_demand = StandardScaler()
                self.scaler_out_of_sample_val_demand = StandardScaler()
                self.scaler_out_of_sample_test_demand = StandardScaler()
            elif self.demand_normalization == 'no_normalization':
                self.scaler_demand = None
                self.scaler_out_of_sample_val_demand = None
                self.scaler_out_of_sample_test_demand = None
            else:
                raise ValueError('demand_normalization must be either "minmax", "standard", or "no_normalization"')
            
            # If demand data used for lag feautures is normalized differently
            if self.lag_demand_normalization != self.demand_normalization:
                if self.lag_demand_normalization == 'minmax':
                    self.scaler_demand_lag = MinMaxScaler()
                    self.scaler_out_of_sample_val_demand_lag = MinMaxScaler()
                    self.scaler_out_of_sample_test_demand_lag = MinMaxScaler()
                elif self.lag_demand_normalization == 'standard':
                    self.scaler_demand_lag = StandardScaler()
                    self.scaler_out_of_sample_val_demand_lag = StandardScaler()
                    self.scaler_out_of_sample_test_demand_lag = StandardScaler()
                elif self.lag_demand_normalization == 'no_normalization':
                    self.scaler_demand_lag = None
                    self.scaler_out_of_sample_val_demand_lag = None
                    self.scaler_out_of_sample_test_demand_lag = None
                else:
                    raise ValueError('lag_demand_normalization must be either "minmax", "standard", or "no_normalization"')
            
            self.scaler_SKU_features = StandardScaler() if self.SKU_features is not None else None # only one since out of sample uses the same fit on known  skus
            self.scaler_time_features = StandardScaler() # only one since time-features are shared between in-sample and out-of-sample SKUs
            self.scaler_time_SKU_features= [StandardScaler() for _ in range(self.num_time_SKU_features)]
            self.scaler_out_of_sample_val_SKU_features = [StandardScaler() for _ in range(self.num_time_SKU_features)]
            self.scaler_out_of_sample_test_SKU_features = [StandardScaler() for _ in range(self.num_time_SKU_features)]

            if initial_normalization:

                logging.info("--Normalizing demand")

                # Prepare data for lag_demand
                self.demand_lag = self.demand.copy() # store original demand values for lag demand

                # Normalize demand targets
                if self.demand_normalization != 'no_normalization':
                    # Normalizing per SKU on time dimension
                    self.scaler_demand.fit(self.demand[:self.train_index_end+1])
                    transformed_demand = self.scaler_demand.transform(self.demand)
                    self.demand.iloc[:,:] = transformed_demand

                # Set unit size for demand targets
                if self.demand_unit_size != None:
                    self.demand = np.round(self.demand, self.demand_unit_size)

                # If separate normalization for lag demand, normalize it
                if self.lag_demand_normalization != self.demand_normalization:
                    if self.lag_demand_normalization != 'no_normalization':
                        self.demand_lag = self.demand.copy() # if lag demand shall be normalized, build on the normalized demand (to account for slight variations due to rounding)
                        self.scaler_demand_lag.fit(self.demand_lag[:self.train_index_end+1])
                        transformed_demand_lag = self.scaler_demand_lag.transform(self.demand_lag)
                        self.demand_lag.iloc[:,:] = transformed_demand_lag
                
                # If lag demand shall be normalized the same way, then copy the normalized demand
                else:
                    self.demand_lag = self.demand.copy()

                if self.SKU_features is not None:
                    logging.info("--Normalizing SKU features")
                    # Normalizing across SKUs, no time dimension present
                    continuous_features = [col for col in self.SKU_features.columns if not self.is_one_hot(self.SKU_features[col])] if ignore_one_hot else self.SKU_features.columns
                    self.SKU_features_to_fit = continuous_features
                    if len(continuous_features) > 0:
                        self.scaler_SKU_features.fit(self.SKU_features[continuous_features]) # SKU features are already calculated based on training index
                        transformed_SKU_features = self.scaler_SKU_features.transform(self.SKU_features[continuous_features])
                        self.SKU_features[continuous_features] = transformed_SKU_features

                logging.info("--Normalizing time features")
                # Normalizting time features (no SKU dimension)
                continuous_features = [col for col in self.time_features.columns if not self.is_one_hot(self.time_features[col])] if ignore_one_hot else self.time_features.columns
                self.time_features_to_fit = continuous_features
                if len(continuous_features) > 0:
                    self.scaler_time_features.fit(self.time_features.loc[:, continuous_features].iloc[:self.train_index_end+1, :])
                    transformed_time_features = self.scaler_time_features.transform(self.time_features.loc[:,continuous_features])
                    self.time_features.loc[:,continuous_features] = transformed_time_features

                logging.info("--Normalizing time-SKU features")
                # Normalize time-SKU features (double-indexed)
                self.time_SKU_features_to_fit = dict()
                for i, feature in enumerate(self.time_SKU_features.columns.get_level_values(0).unique()):
                    # Select all columns corresponding to the current feature in level 0
                    feature_df = self.time_SKU_features.xs(key=feature, axis=1, level=0)

                    should_scale = not ignore_one_hot or (not self.is_one_hot_across_skus(feature_df) if ignore_one_hot else False)
                    self.time_SKU_features_to_fit.update({feature: should_scale})
                    if should_scale:
                        self.scaler_time_SKU_features[i].fit(feature_df[:self.train_index_end+1])
                        transformed_feature_df = self.scaler_time_SKU_features[i].transform(feature_df)
                        self.time_SKU_features.loc[:, (feature, slice(None))] = transformed_feature_df

            
                self.normalized_in_sample_SKUs = True

            else:
                raise NotImplementedError('Training data can only normalized during initialization - later normlization not implemented yet')

    def normalize_demand_and_features_out_of_sample(self,
        normalize: bool = True,
        ignore_one_hot: bool = True,
        initial_normalization = False # Flag if it is set before having added lag features
        ):

        """
        Normalize features using a standard scaler. If ignore_one_hot is true, one-hot encoded features are not normalized.
        """

        if self.out_of_sample_val_SKUs is None and self.out_of_sample_test_SKUs is None:
            return

        if normalize:

            if self.normalized_out_of_sample_SKUs:
                raise ValueError('Features already normalized')

            if initial_normalization:

                logging.info("--Normalizing demand")

                # Prepare data for lag_demand

                self.demand_lag_out_of_sample_test = self.demand_out_of_sample_test.copy()
                self.demand_lag_out_of_sample_val = self.demand_out_of_sample_val.copy()

                # Normalize demand targets
                if self.demand_normalization != 'no_normalization':
                    # Normalizing per SKU on time dimension

                    self.scaler_out_of_sample_test_demand.fit(self.demand_out_of_sample_test[:self.train_index_end+1])
                    transformed_demand = self.scaler_out_of_sample_test_demand.transform(self.demand_out_of_sample_test)
                    self.demand_out_of_sample_test.iloc[:,:] = transformed_demand

                    self.scaler_out_of_sample_val_demand.fit(self.demand_out_of_sample_val[:self.train_index_end+1])
                    transformed_demand = self.scaler_out_of_sample_val_demand.transform(self.demand_out_of_sample_val)
                    self.demand_out_of_sample_val.iloc[:,:] = transformed_demand
                
                # Set unit size for demand targets
                if self.demand_unit_size != None:
                    self.demand_out_of_sample_test = np.round(self.demand_out_of_sample_test, self.demand_unit_size)
                    self.demand_out_of_sample_val = np.round(self.demand_out_of_sample_val, self.demand_unit_size)

                # If separate normalization for lag demand, normalize it
                if self.lag_demand_normalization != self.demand_normalization:
                    if self.lag_demand_normalization != 'no_normalization':
                        
                        self.demand_lag_out_of_sample_test = self.demand_out_of_sample_test.copy()
                        
                        self.demand_lag_out_of_sample_test = self.demand_out_of_sample_test.copy()
                        self.scaler_out_of_sample_test_demand_lag.fit(self.demand_lag_out_of_sample_test[:self.train_index_end+1])
                        transformed_demand_lag = self.scaler_out_of_sample_test_demand_lag.transform(self.demand_lag_out_of_sample_test)
                        self.demand_lag_out_of_sample_test.iloc[:,:] = transformed_demand_lag
                        
                        self.demand_lag_out_of_sample_val = self.demand_out_of_sample_val.copy()
                        self.demand_lag_out_of_sample_val = self.demand_out_of_sample_val.copy()
                        self.scaler_out_of_sample_val_demand_lag.fit(self.demand_lag_out_of_sample_val[:self.train_index_end+1])
                        transformed_demand_lag = self.scaler_out_of_sample_val_demand_lag.transform(self.demand_lag_out_of_sample_val)
                        self.demand_lag_out_of_sample_val.iloc[:,:] = transformed_demand_lag
                    
                # If lag demand shall be normalized the same way, then copy the normalized demand
                else:
                    self.demand_lag_out_of_sample_test = self.demand_out_of_sample_test.copy()
                    self.demand_lag_out_of_sample_val = self.demand_out_of_sample_val.copy()
                    
                if self.SKU_features is not None:
                    logging.info("--Normalizing SKU features")
                    # Normalizing across SKUs, no time dimension present
                    continuous_features = self.SKU_features_to_fit
                    if len(continuous_features) > 0:
                        transformed_SKU_features = self.scaler_SKU_features.transform(self.SKU_features_out_of_sample_test[continuous_features])
                        self.SKU_features_out_of_sample_test[continuous_features] = transformed_SKU_features
                    
                    if len(continuous_features) > 0:
                        transformed_SKU_features = self.scaler_SKU_features.transform(self.SKU_features_out_of_sample_val[continuous_features])
                        self.SKU_features_out_of_sample_val[continuous_features] = transformed_SKU_features

                logging.info("--Normalizing time-SKU features")
                
                for i, feature in enumerate(self.time_SKU_features_out_of_sample_test.columns.get_level_values(0).unique()):
                    # Select all columns corresponding to the current feature in level 0
                    feature_df = self.time_SKU_features_out_of_sample_test.xs(key=feature, axis=1, level=0)

                    should_scale = self.time_SKU_features_to_fit[feature]
                    if should_scale:
                        self.scaler_out_of_sample_test_SKU_features[i].fit(feature_df[:self.train_index_end+1])
                        transformed_feature_df = self.scaler_out_of_sample_test_SKU_features[i].transform(feature_df)
                        self.time_SKU_features_out_of_sample_test.loc[:, (feature, slice(None))] = transformed_feature_df
                
                for i, feature in enumerate(self.time_SKU_features_out_of_sample_val.columns.get_level_values(0).unique()):
                    # Select all columns corresponding to the current feature in level 0
                    feature_df = self.time_SKU_features_out_of_sample_val.xs(key=feature, axis=1, level=0)

                    should_scale = self.time_SKU_features_to_fit[feature]
                    if should_scale:
                        self.scaler_out_of_sample_val_SKU_features[i].fit(feature_df[:self.train_index_end+1])
                        transformed_feature_df = self.scaler_out_of_sample_val_SKU_features[i].transform(feature_df)
                        self.time_SKU_features_out_of_sample_val.loc[:, (feature, slice(None))] = transformed_feature_df
            
                self.normalized_out_of_sample_SKUs = True

            else:
                raise NotImplementedError('Training data can only normalized during initialization - later normlization not implemented yet')

    def update_lag_features(self,
        lag_window: int,
        ):

        """ Update lag window parameters for dataloader object that is already initialized """

        raise NotImplementedError('Not implemented yet')

        # Problem: updating lag_features naively would shorten the dataset each time it is called

    def get_time_SKU_idx(self, idx):

        """ get time and SKU index by index, depending on the dataset type (train, val, test) """

        # print("mode of dataloader when getting item:", self.dataset_type)

        if self.dataset_type == "train":

            if self.meta_learn_units:

                if idx >= len(self.sku_time_index):
                    raise IndexError(f'index {idx} out of range{len(self.sku_time_index)}')
                idx_sku, idx_time, = self.sku_time_index[idx]
                idx_skus = [idx_sku]

            else:
                if idx+self.train_index_start > self.train_index_end:
                    raise IndexError(f'index {idx} out of range{self.train_index_end-self.train_index_start}')
                idx_skus = self.train_SKUs_indices
                idx_time = idx
            idx_time += self.train_index_start

        elif self.dataset_type == "val":
            idx_time = idx + self.val_index_start
            if self.return_SKU_type == "in_sample":
                if self.in_sample_val_test_SKUs is not None:
                    idx_skus = self.in_sample_val_test_SKUs_indices
                else:
                    idx_skus = self.train_SKUs_indices
            elif self.return_SKU_type == "out_of_sample_val":
                idx_skus = self.out_of_sample_val_SKUs_indices
            elif self.return_SKU_type == "out_of_sample_test":
                idx_skus = self.out_of_sample_test_SKUs_indices
            else:
                raise ValueError('return_SKU_type not set')
            
            if idx >= self.test_index_start:
                raise IndexError(f'index{idx} out of range{self.test_index_start}')
        
        elif self.dataset_type == "test":
            idx_time = idx + self.test_index_start
            if self.return_SKU_type == "in_sample":
                if self.in_sample_val_test_SKUs is not None:
                    idx_skus = self.in_sample_val_test_SKUs_indices
                else:
                    idx_skus = self.train_SKUs_indices
            elif self.return_SKU_type == "out_of_sample_val":
                idx_skus = self.out_of_sample_val_SKUs_indices
            elif self.return_SKU_type == "out_of_sample_test":
                idx_skus = self.out_of_sample_test_SKUs_indices
            else:
                raise ValueError('return_SKU_type not set')
            
            if idx >= len(self.demand):
                raise IndexError(f'index{idx} out of range{len(self.demand)}')
        else:
            raise ValueError('dataset_type not set')

        return idx_time, idx_skus

    def __getitem__(self, idx: int):

        """ get item by index, depending on the dataset type (train, val, test)"""

        if self.dataset_type != "train":

            time_features = self.time_features # time features independent of SKU

            if self.return_SKU_type == "in_sample":
                demand = self.demand
                demand_lag = self.demand_lag
                SKU_features = self.SKU_features
                time_SKU_features = self.time_SKU_features
                mask = self.mask
                len_SKUs = len(self.train_SKUs)
            elif self.return_SKU_type == "out_of_sample_val":
                demand = self.demand_out_of_sample_val
                demand_lag = self.demand_lag_out_of_sample_val
                SKU_features = self.SKU_features_out_of_sample_val
                time_SKU_features = self.time_SKU_features_out_of_sample_val
                mask = self.mask_out_of_sample_val
                len_SKUs = len(self.out_of_sample_val_SKUs)
            elif self.return_SKU_type == "out_of_sample_test":
                demand = self.demand_out_of_sample_test
                demand_lag = self.demand_lag_out_of_sample_test
                SKU_features = self.SKU_features_out_of_sample_test
                time_SKU_features = self.time_SKU_features_out_of_sample_test
                mask = self.mask_out_of_sample_test
                len_SKUs = len(self.out_of_sample_test_SKUs)
        
        else:
            demand = self.demand
            demand_lag = self.demand_lag
            SKU_features = self.SKU_features
            time_SKU_features = self.time_SKU_features
            mask = self.mask
            len_SKUs = len(self.train_SKUs)
        
        time_features = self.time_features # time features independent of SKU


        lag_window = self.lag_window_params["lag_window"]
        include_y = self.lag_window_params["include_y"]

        idx_time, idx_skus = self.get_time_SKU_idx(idx)
        num_skus = len(idx_skus)

        demand = demand[idx_time, idx_skus]
        
        item = np.empty((1,lag_window+1, self.num_features, num_skus))

        if include_y:
            assert idx_time-1-lag_window >= 0
            lag_demand = demand_lag[idx_time-1-lag_window:idx_time, idx_skus] # need to use t-1 to get the lag  

        if self.provide_additional_target:
            additional_target = demand_lag[idx_time-lag_window:idx_time+1, idx_skus] # provide target without lag
            additional_target[-1, :] = 0 # The transformer cannot see the last target --> this is to be predicted

        if self.SKU_features is not None:
            SKU_features = SKU_features[idx_skus].transpose()
            len_SKU_features = len(SKU_features)
            SKU_features = np.expand_dims(SKU_features, axis=0)
            SKU_features = np.repeat(SKU_features, repeats=lag_window+1, axis=0) 
        else:
            len_SKU_features = 0

        time_features = time_features[idx_time-lag_window:idx_time+1]
        len_time_features = time_features.shape[1]
        time_features = np.expand_dims(time_features, axis=-1)
        time_features = np.repeat(time_features, num_skus, axis=-1)

        num_time_SKU_features_without_lag_demand = self.num_time_SKU_features-include_y-self.include_non_available-self.provide_additional_target

        for i, idx_sku in enumerate(idx_skus):

            SKU_indices = [len_SKUs*i+idx_sku for i in range(num_time_SKU_features_without_lag_demand)]
            time_SKU_features_sku = time_SKU_features[idx_time-lag_window:idx_time+1, SKU_indices]
            item[:,:,(len_SKU_features+len_time_features):(len_SKU_features+len_time_features+num_time_SKU_features_without_lag_demand),i] = np.expand_dims(time_SKU_features_sku, axis=0)
        
        item[:,:,:len_SKU_features,:] =  np.expand_dims(SKU_features, axis=0)
        
        item[:,:,len_SKU_features:(len_SKU_features+len_time_features),:] = np.expand_dims(time_features, axis=0)

        extra_info = sum([self.include_non_available, include_y, self.provide_additional_target])
        additional_info = np.empty((1,lag_window+1, extra_info, num_skus))

        current_index = 0

        if self.include_non_available:

            additional_info[:, :, current_index, :] = np.expand_dims(mask[idx_time-lag_window:idx_time+1, idx_skus], axis=0)
            current_index += 1

        if include_y:
        
            additional_info[:, :, current_index, :] = np.expand_dims(lag_demand, axis=0)
            current_index += 1
    
        if self.provide_additional_target:
            
            additional_info[:, :, current_index, :] = np.expand_dims(additional_target, axis=0)

        item[:,:,-extra_info:,:] = additional_info

        if self.dataset_type == "train":
            if self.permutate_inputs:
                start_index_to_permutate = len_SKU_features
                end_index_to_permutate = item.shape[2]
                if self.provide_additional_target:
                    end_index_to_permutate -= 1 # target shall always be at the end
                indices_for_permutation = np.arange(start_index_to_permutate, end_index_to_permutate)
                print("original order:", indices_for_permutation)
                indices = np.random.permutation(indices_for_permutation)
                item[:,:,(start_index_to_permutate):end_index_to_permutate,:] = item[:,:,(indices),:]
                print("new order:", indices)

        if self.meta_learn_units:
            if self.dataset_type == "train":
                if item.shape[-1] == 1:
                    item = item.squeeze(-1)
                else:
                    raise ValueError('SKU as batch, but item has more than one SKU dimension')
        else:
            if item.shape[-1] != 1:
                raise ValueError('Num_units dimension must be 1 if not meta-learning')
            else:
                item = item.squeeze(-1)

        if item.shape[0] != 1:
            raise ValueError('Batch dimension must be 1')
        item = item.squeeze(0) # remove batch dimension
 
        return item, demand

    def __len__(self):
        return len(self.demand)
    
    @property
    def X_shape(self):

        if self.meta_learn_units:
            return (len(self.time_features), self.lag_window_params["lag_window"]+1, self.num_features)
        else:
            return (len(self.time_features), self.lag_window_params["lag_window"]+1, self.num_features) # check if there will be a difference.
    
    @property
    def Y_shape(self):
        return (len(self.time_features), self.num_units) # using num of val_test SKUs

    @property
    def len_train(self):
        if self.meta_learn_units:
            return len(self.sku_time_index) # sku_time_index contains only timesteps that are in the training set and skus in the training set.
        else:
            return self.len_train_time

    @property
    def len_val(self):
        if self.val_index_start is None:
            raise ValueError('no validation set defined')
        return self.test_index_start-self.val_index_start # validating and testing is always along the time demension (units are a separate dimension)

    @property
    def len_test(self):
        if self.test_index_start is None:
            raise ValueError('no test set defined')
        return len(self.demand)-self.test_index_start # validating and testing is always along the time demension (units are a separate dimension)

    def get_all_X(self,
                dataset_type: str = 'train' # can be 'train', 'val', 'test', 'all'
                ): 

        logging.info("Retrieving all X data")

        """
        Returns the entire features dataset.
        Return either the train, val, test, or all data.
        """

        print("mode of dataloader when getting item:", self.dataset_type)

        if dataset_type == 'train':

            X = np.empty((self.len_train, self.lag_window_params["lag_window"]+1, self.num_features))

            for i in range(self.len_train):
                X[i], _ = self[i]

        elif dataset_type == 'val':
            raise NotImplementedError('Not implemented yet')
        elif dataset_type == 'test':
            raise NotImplementedError('Not implemented yet')
        elif dataset_type == 'all':
            raise NotImplementedError('Not implemented yet')
        else:
            raise ValueError('dataset_type not recognized')

        return X

    def get_all_Y(self,
                dataset_type: str = 'train' # can be 'train', 'val', 'test', 'all'
                ): 

        """
        Returns the entire target dataset.
        Return either the train, val, test, or all data.
        """

        print("mode of dataloader when getting item:", self.dataset_type)
        if dataset_type == 'train':
            return self.demand[self.train_index_start:self.val_index_start, self.train_SKUs_indices]
        elif dataset_type == 'val':
            raise NotImplementedError('Not implemented yet')
        elif dataset_type == 'test':
            raise NotImplementedError('Not implemented yet')
        elif dataset_type == 'all':
            raise NotImplementedError('Not implemented yet')
        else:
            raise ValueError('dataset_type not recognized')

    @staticmethod
    def is_one_hot(column):
        return set(column.unique()) <= {0, 1}

    @staticmethod
    def is_one_hot_across_skus(feature_df):
        """
        Check if the set of unique values in a feature across all SKU_ids is {0, 1}.
        feature_df: DataFrame slice for a specific feature with SKU_ids as columns.
        """
        flattened_values = feature_df.values.flatten()
        
        # Check if the unique values in this flattened array are exactly {0, 1}
        unique_values = set(flattened_values)
        
        return unique_values <= {0, 1}

    @staticmethod
    def save_indices(df):
        """
        Saves the row and column indices of a DataFrame.
        """
        return {
            'rows': df.index,
            'columns': df.columns
        }
    
    def set_return_sku(
        self,
        sku_type: str = "in_sample" # if train, then the dataloader returns in-sample SKUs, other option "out_of_sample_val" or "out_of_sample_test"
        ):
       
        """
        Set the SKU type to be returned in the __getitem__ method.
        """

        self.return_SKU_type = sku_type
