## TODO

## Accepts a dataset name and model id, predict and produces graph + error metrics


import matplotlib.pyplot as plt
import mlflow
import mlflow.data
import numpy as np
import pandas as pd
from mlflow.client import MlflowClient
from mlflow.data.pandas_dataset import PandasDataset
from utilsforecast.plotting import plot_series

from neuralforecast.core import NeuralForecast
from neuralforecast.models import NBEATSx
from neuralforecast.utils import AirPassengersDF
from neuralforecast.losses.pytorch import MAE

import matplotlib.pyplot as plt
Y_plot = Y_hat_df[Y_hat_df['unique_id']=='Dated']
cutoffs = Y_hat_df['cutoff'].unique()[::horizon]
Y_plot = Y_plot[Y_hat_df['cutoff'].isin(cutoffs)]

plt.figure(figsize=(20,5))
plt.plot(Y_plot['ds'], Y_plot['y'], label='True')
for model in models:
    plt.plot(Y_plot['ds'], Y_plot[f'{model}'], label=f'{model}')
plt.xlabel('Datestamp')
plt.ylabel('OT')
plt.grid()
plt.legend()

from neuralforecast.losses.numpy import mse, mae, mape

for model in models:
    mae_model = mae(Y_hat_df['y'], Y_hat_df[f'{model}'])
    mse_model = mse(Y_hat_df['y'], Y_hat_df[f'{model}'])
    mape_model = mape(Y_hat_df['y'], Y_hat_df[f'{model}'])
    print(f'{model} horizon {horizon} - MAE: {mae_model:.3f}')
    print(f'{model} horizon {horizon} - MSE: {mse_model:.3f}')
    print(f'{model} horizon {horizon} - MAPE: {mape_model:.3f}')