from mlforecast import MLForecast
from mlforecast.lag_transforms import RollingMean
from sklearn.base import BaseEstimator
from utilsforecast.data import generate_series

class WindowAverage(BaseEstimator):
    def __init__(self, window_size):
        self.window_size = window_size

    def fit(self, X, y):
        return self

    def predict(self, X):
        return X[f'rolling_mean_lag1_window_size{self.window_size}']

freq = 'H'
window_size = 13
horizon = 2 * window_size
data = generate_series(1, freq=freq)
fcst = MLForecast(
    models=[WindowAverage(window_size)],
    freq=freq,
    lag_transforms={1: [RollingMean(window_size)]}
)
fcst.fit(data)
fcst.predict(horizon)