import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datasetsforecast.long_horizon import LongHorizon
from neuralforecast.core import NeuralForecast
from neuralforecast.models import NHITS, PatchTST, iTransformer, TSMixer
from utilsforecast.losses import mae, mse
from utilsforecast.evaluation import evaluate


def load_data(name):
    # LongHorizon: https://nhits-experiments.s3.amazonaws.com/datasets.zip
    '''
    unique_id: the unique identifier of the time series.
    Y_df: training data.
    ds: the datestamp of the forecast for each row.
    y: the actual value of the target variable.
    '''
    '''
    # https://github.com/Nixtla/datasetsforecast/blob/main/datasetsforecast/long_horizon.py
    y_df = pd.read_csv(f'{path}/{name}/{kind}/df_y.csv')
    y_df = y_df.sort_values(['unique_id', 'ds'], ignore_index=True)
    y_df = y_df[['unique_id', 'ds', 'y']]
    X_df = pd.read_csv(f'{path}/{name}/{kind}/df_x.csv')
    X_df = y_df.drop('y', axis=1).merge(X_df, how='left', on=['ds'])
    S_df = None
    '''

    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


horizon = 96

# train
Y_df, val_size, test_size, freq = load_data('ettm2')

models = [
    iTransformer(h=horizon, input_size=3*horizon, n_series=1, max_steps=1000, early_stop_patience_steps=3),
] # n_series:features num

nf = NeuralForecast(models=models, freq=freq)
# nf.fit(df=Y_df, val_size=val_size) # DataFrame with columns [`unique_id`, `ds`, `y`], target_col : str (default='y').
nf_preds = nf.cross_validation(df=Y_df, val_size=val_size, test_size=test_size, n_windows=None)
nf_preds = nf_preds.reset_index()
nf.save(path='./checkpoints/test_run/',
    model_index=None, 
    overwrite=True,
    save_dataset=False)

evaluation = evaluate(df=nf_preds, metrics=[mae, mse], models=['iTransformer'])
evaluation.to_csv(f'ettm1_results.csv', index=False, header=True)


# 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()
'''