import pickle
from neuralforecast import NeuralForecast
from modules.transform import calendarFeatures
import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import CSVLogger

def main(artifacts_path, variate='uni', horizon_len=5):
    with open(os.path.join(artifacts_path, 'transformations.pkl'), 'rb') as fp:
        transformations_loaded = pickle.load(fp)
    
    nf2 = NeuralForecast.load(path=os.path.join(artifacts_path, 'model'))
    nf2.freq = 'B'

    training_data_path = os.path.join(artifacts_path, 'train_data.csv')
    Y_df = pd.read_csv(training_data_path)
    Y_df['ds'] = pd.to_datetime(Y_df['ds'])

    futr_df = nf2.make_future_dataframe()

    if variate=='univariate':
        futr_df = calendarFeatures(futr_df)
        Y_hat = nf2.predict(futr_df=futr_df).reset_index()
        plot_data = Y_df.copy()
        plot_data = plot_data[plot_data['ds'] >= (plot_data['ds'].max() - pd.Timedelta(days=horizon_len * 2))]

        sns.set_theme(style="whitegrid")
        fig, ax = plt.subplots(1, 1, figsize=(20, 7))
        sns.lineplot(data=plot_data, x='ds', y='y', label='Actual', ax=ax, color='blue', linewidth=2)
        for col in Y_hat.columns:
            if col not in ['ds', 'unique_id']:
                sns.lineplot(data=Y_hat, x='ds', y=col, label=col, ax=ax, linewidth=2)
        
        ax.set_title('Forecasting Result', fontsize=22)
        ax.set_ylabel('Y', fontsize=20)
        ax.set_xlabel('Date', fontsize=20)
        ax.legend(prop={'size': 15})
        ax.grid(True)

        plot_path = os.path.join(artifacts_path, 'forecast_plot.jpg')
        plt.savefig(plot_path, format='jpg', dpi=300)

    if variate=='multivariate':
        Y_hat = nf2.predict(futr_df=futr_df).reset_index()
        plot_data = Y_df.copy()
        plot_data = plot_data[plot_data['ds'] >= (plot_data['ds'].max() - pd.Timedelta(days=horizon_len * 2))]

        sns.set_theme(style="whitegrid")
        unique_ids = Y_hat['unique_id'].unique()
        fig, axes = plt.subplots(len(unique_ids), 1, figsize=(20, 7 * len(unique_ids)), sharex=True)
        if len(unique_ids) == 1:
            axes = [axes]

        for ax, uid in zip(axes, unique_ids):
            actual_data = plot_data[plot_data['unique_id'] == uid]
            forecast_data = Y_hat[Y_hat['unique_id'] == uid]

            sns.lineplot(data=actual_data, x='ds', y='y', label='Actual', ax=ax, color='blue', linewidth=2)
            for col in Y_hat.columns:
                if col not in ['ds', 'unique_id']:
                    sns.lineplot(data=forecast_data, x='ds', y=col, label=col, ax=ax, linewidth=2)

            ax.set_title(f'Forecasting Result for {uid}', fontsize=22)
            ax.set_ylabel('Y', fontsize=20)
            ax.set_xlabel('Date', fontsize=20)
            ax.legend(prop={'size': 15})
            ax.grid(True)

        plt.tight_layout()
        plot_path = os.path.join(artifacts_path, 'forecast_plot.jpg')
        plt.savefig(plot_path, format='jpg', dpi=300)
    Y_hat.to_csv(os.path.join(artifacts_path, 'prediction_results.csv'))

if __name__ == "__main__":
    directory = Path(__file__).parent.absolute()
    run_id = 'crude_oil_d2eea033-7bf9-437d-809c-d18f574a982d'  # This should be inputted (unique id for every run)
    artifacts_path = os.path.join(directory, 'artifacts', run_id)

    log_dir = os.path.join(artifacts_path, 'prediction_logs')
    Trainer.default_root_dir = log_dir
    
    logger = CSVLogger(save_dir=log_dir, name='forecast_logs')

    main(artifacts_path=artifacts_path, variate='univariate', horizon_len=5)