"""THE ARIMA Model and its settings """



import pandas as pd
import matplotlib.pyplot as plt
import cftime
import yaml


# Load configuration from YAML file
def load_config(config_path):
    with open(config_path, 'r') as file:
        return yaml.safe_load(file)


# Function to convert cftime to datetime
def convert_cftime_to_datetime(index):
    return [pd.Timestamp(date.year, date.month, date.day, date.hour, date.minute, date.second) for date in index]


# Function to plot time series data
def plot_time_series(df_train, df_test, config):
    # Convert the index of both DataFrames
    df_train.index = convert_cftime_to_datetime(df_train.index)
    df_test.index = convert_cftime_to_datetime(df_test.index)

    # Plot each column individually
    num_plots = len(df_train.columns.union(df_test.columns))
    fig, axes = plt.subplots(nrows=num_plots, ncols=1, figsize=config['figure_size'], sharex=True)

    if num_plots == 1:
        axes = [axes]  # Ensure axes is iterable if only one subplot

    for i, column in enumerate(df_train.columns.union(df_test.columns)):
        ax = axes[i]

        if column in df_train.columns:
            ax.plot(df_train.index, df_train[column], label=f'Train {column}', color='blue')

        if column in df_test.columns:
            ax.plot(df_test.index, df_test[column], label=f'Test {column}', color='red')

        ax.set_title(f'Time Series for {column}')
        ax.set_xlabel('Date')
        ax.set_ylabel('Value')
        ax.legend()

    plt.tight_layout()
    plt.show()


# Main function to run the script
def main(config_path, df_train, df_test):
    config = load_config(config_path)
    plot_time_series(df_train, df_test, config)

# Example usage (uncomment and set paths as needed):
# if __name__ == "__main__":
#     config_path = 'path/to/config.yml'
#     # df_train and df_test should be defined or loaded before this
#     main(config_path, df_train, df_test)
