from time import time
import ray
import pandas as pd
#from neuralforecast.data.datasets.m5 import M5, M5Evaluation
from statsforecast import StatsForecast
from statsforecast.models import ETS

ray.init(address="auto")

Y_df = pd.read_parquet('<s3://m5-benchmarks/data/train/target.parquet>')
Y_df = Y_df.rename(columns={
    'item_id': 'unique_id', 
    'timestamp': 'ds', 
    'demand': 'y'
})
Y_df['ds'] = pd.to_datetime(Y_df['ds'])

Y_df = Y_df[Y_df.unique_id == "FOODS_1_001_CA_1"]

constant = 10
Y_df['y'] += constant

fcst = StatsForecast(
    df=Y_df, 
    models=[ETS(season_length=7, model='ZNA')], 
    freq='D'
)

Y_hat = fcst.forecast(28)