# AUTOGENERATED! DO NOT EDIT! File to edit: model.data.ipynb (unless otherwise specified).

__all__ = ['train_val_datesplit', 'add_lagged_features', 'add_rolling_features']

# Cell

def train_val_datesplit(df, train_start, train_end, val_end, obs_col='obs'):
    """Split into training and validation sets based on dates"""

    train_df = df[(df.index >= train_start) & (df.index <= train_end)]
    val_df = df[(df.index > train_end) & (df.index <= val_end)]

    X_train = train_df.drop([obs_col], axis=1)
    y_train = train_df[[obs_col]]

    X_val = val_df.drop([obs_col], axis=1)
    y_val = val_df[[obs_col]]

    return (X_train, y_train, X_val, y_val)

# Cell

def add_lagged_features(df, features, lags):
    for feature in features:
        for lag in lags:
            df[f'{feature}_lagged_{lag}_timestamp'] = df[feature].shift(lag)
    return df.dropna()

# Cell

def add_rolling_features(df, features, agg_name, agg_days):
    for feature in features:
        for days in agg_days:
            if agg_name == 'sum':
                df[f'{feature}_rolling_sum_{days}_days'] = df[feature].rolling(f'{days}D').sum()
            if agg_name == 'mean':
                df[f'{feature}_rolling_mean_{days}_days'] = df[feature].rolling(f'{days}D').mean()

    return df.dropna()