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

__all__ = ['log', 'add_lags']

# Cell

import logging
from typing import Union

import numpy as np
import pandas as pd

log = logging.getLogger(__name__)

# Cell


def add_lags(data: Union[np.array, pd.DataFrame], n_lags: int = 1) -> Union[np.array, pd.DataFrame]:
    """Take in a numpy array or pandas df and add n_lags for each column.
    """
    data_orig = np.copy(data)
    if isinstance(data, np.ndarray):
        for n_lag in range(1, n_lags + 1):
            data = np.concatenate((data, np.roll(data_orig, n_lag, axis=0)), axis=1)
        data = data[n_lags:]
    elif isinstance(data, pd.DataFrame):
        colnames_new = [f"{col}_lag{n_lag}".replace('_lag0', '') for n_lag in range(n_lags+1) for col in data.columns]
        data = pd.concat([data.shift(n_lag) for n_lag in range(n_lags + 1)], axis=1)
        data.columns = colnames_new
    log.debug(f'... (add_lags) n_lags = {n_lags} arr_orig.shape = {data_orig.shape}  arr.shape = {data.shape}')
    return data

