scaler = TemporalNorm(scaler_type='standard', dim=1)

import numpy as np
from coreforecast.grouped_array import GroupedArray as CoreGroupedArray
from mlforecast import MLForecast
from mlforecast.grouped_array import GroupedArray
from mlforecast.lag_transforms import RollingMean
from mlforecast.target_transforms import _BaseGroupedArrayTargetTransform
from sklearn.linear_model import LinearRegression
from utilsforecast.data import generate_series

class WindowStandardScaler(_BaseGroupedArrayTargetTransform):
    def __init__(self, window_size: int):
        self.window_size = window_size

    def update(self, ga: GroupedArray) -> GroupedArray:
        raise NotImplementedError

    def take(self, idxs: np.ndarray) -> 'WindowStandardScaler':
        raise NotImplementedError

    def update(self, ga: GroupedArray) -> GroupedArray:
        raise NotImplementedError

    def fit_transform(self, ga: GroupedArray) -> GroupedArray:
        rm = RollingMean(self.window_size, min_samples=1)
        rm._set_core_tfm(lag=0)
        core_ga = CoreGroupedArray(ga.data, ga.indptr)
        scales = rm.transform(core_ga)
        self.last_scales_ = CoreGroupedArray(scales, ga.indptr)._tail(1)
        return GroupedArray(ga.data / scales, ga.indptr)

    def inverse_transform(self, ga: GroupedArray) -> GroupedArray:
        h = ga.data.size // (ga.indptr.size - 1)
        scales = np.repeat(self.last_scales_, h)
        return GroupedArray(ga.data * scales, ga.indptr)

series = generate_series(10, min_length=100)
fcst = MLForecast(
    models=[LinearRegression()],
    freq="D",
    lags=[1, 2],
    target_transforms=[WindowStandardScaler(window_size=7)]
)
fcst.fit(series)
fcst.predict(5)