import pandas as pd
import numpy as np
from datasetsforecast.long_horizon import LongHorizon
from neuralforecast.core import NeuralForecast


def load_data(name):
    if name == "ettm2":
        
        Y_df, X_df, S_df = LongHorizon.load(directory='./ETT-small/', group='ETTm2')
        Y_df = Y_df[Y_df['unique_id'] == 'OT']
        Y_df['ds'] = pd.to_datetime(Y_df['ds'])
        val_size = 11520
        test_size = 11520
        freq = '15T'

    return Y_df, val_size, test_size, freq


# infer
Y_df, val_size, test_size, freq = load_data('ettm2')
nf = NeuralForecast.load(path='./checkpoints/test_run/')

Y_hat_df = nf.predict(Y_df).reset_index()#_predict(df: pd.DataFrame, static_cols, futr_exog_cols, models, freq, id_col, time_col, target_col)
print("Y_hat_df: ", Y_hat_df)




'''
futr_df = pd.read_csv('https://datasets-nixtla.s3.amazonaws.com/EPF_FR_BE_futr.csv')
futr_df['ds'] = pd.to_datetime(futr_df['ds'])
Y_hat_df = nf.predict(futr_df=futr_df)
Y_hat_df.head()
'''

'''
from neuralforecast.utils import AirPassengersDF
Y_df = AirPassengersDF # Defined in neuralforecast.utils
Y_df.head()
'''