import pandas as pd
import matplotlib.pyplot as plt

from neuralforecast import NeuralForecast
from neuralforecast.models import LSTM, NHITS, RNN, PatchTST


data = pd.read_csv("terheles_fixed.tsv", sep="\t")
data['ds'] = pd.to_datetime(data['Korrigált időpont'])
data['y'] = data['Hatásos teljesítmény']

data = data[['ds', 'y']]
data['unique_id'] = 1

data = data[data['ds'] < '2019-09-01']
Y_df = data
print(Y_df)


horizon = 4 * 24 * 7 # 7 days

# Try different hyperparmeters to improve accuracy.
models = [
            PatchTST(h=horizon,                    # Forecast horizon
               input_size=2 * horizon,      # Length of input sequence
               max_steps=20,                # Number of steps to train
               scaler_type='standard'),       # Type of scaler to normalize data

            NHITS(h=horizon,                   # Forecast horizon
                input_size=2 * horizon,      # Length of input sequence
                max_steps=100,               # Number of steps to train
                n_freq_downsample=[2, 1, 1]) # Downsampling factors for each stack output
          ]
'''
            LSTM(h=horizon,                    # Forecast horizon
               max_steps=100,                # Number of steps to train
               scaler_type='standard',       # Type of scaler to normalize data
               encoder_hidden_size=64,       # Defines the size of the hidden state of the LSTM
               decoder_hidden_size=64,),     # Defines the number of hidden units of each layer of the MLP decoder
'''


nf = NeuralForecast(models=models, freq='15min')

shorter_Y_df = Y_df[Y_df['ds'] < '2019-08-01']
print("-=======-")
print(len(shorter_Y_df))
print(shorter_Y_df)
nf.fit(df=shorter_Y_df)

Y_hat_df = nf.predict()

print(Y_df)

Y_hat_df = Y_hat_df.reset_index()

fig, ax = plt.subplots(1, 1, figsize = (20, 7))

# plot_df = pd.concat([Y_df, Y_hat_df]).set_index('ds') # Concatenate the train and forecast dataframes
# plot_df[['y', 'LSTM', 'NHITS']].plot(ax=ax, linewidth=2)

plot_Y_df = Y_df[Y_df['ds'] > '2019-07-01']
plot_Y_df = plot_Y_df.set_index('ds')[['y']]
plot_Y_df.plot(ax=ax, linewidth=1)
Y_hat_df.set_index('ds')[['PatchTST', 'NHITS']].plot(ax=ax, linewidth=1)


ax.set_title('AirPassengers Forecast', fontsize=22)
ax.set_ylabel('Monthly Passengers', fontsize=20)
ax.set_xlabel('Timestamp [t]', fontsize=20)
ax.legend(prop={'size': 15})
ax.grid()

plt.savefig("neuralforecast.pdf")