ARIMA Forecast

# Dependency Installation
!pip install pandas numpy matplotlib seaborn statsmodels scipy scikit-learn

# Library Imports
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import statsmodels.api as sm
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.stattools import adfuller
from sklearn.metrics import mean_squared_error
from collections import deque
import logging
import time
import random
import math
from statsmodels.tsa.arima_process import arma_generate_sample # Added this import

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Set seaborn style
sns.set_style("whitegrid")

# Set random seed for reproducibility
np.random.seed(42)
random.seed(42)

def create_state(
    start_date: str = '2020-01-01',
    end_date: str = '2023-01-01',
    frequency: str = 'D',
    ar_params: list[float] = None,
    ma_params: list[float] = None,
    n_samples: int = 1000,
    random_seed: int = 42
) -> dict:
    """
    Initializes the global state dictionary with configuration and empty data structures.

    Parameters
    ----------
    start_date : str, optional
        The start date for generating time series data in 'YYYY-MM-DD' format, by default '2020-01-01'.
    end_date : str, optional
        The end date for generating time series data in 'YYYY-MM-DD' format, by default '2023-01-01'.
    frequency : str, optional
        The frequency of the time series data (e.g., 'D' for daily, 'M' for monthly), by default 'D'.
    ar_params : list[float], optional
        AutoRegressive (AR) coefficients for data generation, by default None.
    ma_params : list[float], optional
        Moving Average (MA) coefficients for data generation, by default None.
    n_samples : int, optional
        The number of samples to generate for the time series, by default 1000.
    random_seed : int, optional
        Seed for random number generators to ensure reproducibility, by default 42.

    Returns
    -------
    dict
        An initialized state dictionary containing all configuration and empty data structures.
    """
    logger.info("Initializing application state...")
    if ar_params is None:
        ar_params = [0.7]
    if ma_params is None:
        ma_params = [0.5]

    state = {
        'config': {
            'start_date': start_date,
            'end_date': end_date,
            'frequency': frequency,
            'ar_params': ar_params,
            'ma_params': ma_params,
            'n_samples': n_samples,
            'random_seed': random_seed,
            'test_size': 0.2,
            'retries': 3,
            'backoff_factor': 0.5
        },
        'data': {
            'original': None,
            'train': None,
            'test': None,
            'stationary': None
        },
        'model': {
            'fitted_model': None,
            'arima_order': (0, 0, 0),
            'forecast_steps': 0
        },
        'results': {
            'predictions': None,
            'metrics': {}
        },
        'logs': deque(maxlen=100) # Store recent logs
    }
    logger.debug(f"Initial state created: {state['config']}")
    return state

def create_time_series_data(state: dict) -> dict:
    """
    Generates synthetic time series data based on AR and MA parameters.

    Parameters
    ----------
    state : dict
        The current state dictionary containing configuration for data generation.
        Expected keys: 'start_date', 'end_date', 'frequency', 'ar_params', 'ma_params', 'random_seed'.

    Returns
    -------
    dict
        The updated state dictionary with the generated time series data.
    """
    logger.info("Generating synthetic time series data...")
    config = state['config']

    np.random.seed(config['random_seed'])
    random.seed(config['random_seed'])

    try:
        # Generate dates
        dates = pd.date_range(start=config['start_date'], end=config['end_date'], freq=config['frequency'])
        n_samples = len(dates)

        # Ensure AR and MA polynomials are correctly formatted for arma_generate_sample
        # The polynomial for AR coefficients is 1 - ar1*L - ar2*L^2 - ...
        # The polynomial for MA coefficients is 1 + ma1*L + ma2*L^2 - ...
        ar_coeffs = np.r_[1, -np.array(config['ar_params'])]
        ma_coeffs = np.r_[1, np.array(config['ma_params'])]

        # Generate data using arma_generate_sample
        # We need to ensure the number of samples matches the date range.
        # The original config['n_samples'] might be different if a fixed range is specified.
        # Adjusting n_samples to match the length of the date range.
        simulated_data = arma_generate_sample(
            ar=ar_coeffs, ma=ma_coeffs, nsample=n_samples, scale=0.5, distrvs=np.random.normal
        )

        # Create DataFrame
        time_series_df = pd.DataFrame(simulated_data, index=dates, columns=['value'])
        state['data']['original'] = time_series_df
        logger.info(f"Generated {len(time_series_df)} data points from {dates.min().strftime('%Y-%m-%d')} to {dates.max().strftime('%Y-%m-%d')}.")
        state['logs'].append(f"Time series data generated with {len(time_series_df)} points.")
    except Exception as e:
        logger.error(f"Error generating time series data: {e}")
        state['logs'].append(f"ERROR: Failed to generate time series data - {e}")
    return state

def plot_time_series(state: dict, title: str = 'Time Series Data', ylabel: str = 'Value') -> dict:
    """
    Plots the original time series data.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting 'original' time series data under state['data'].
    title : str, optional
        Title of the plot, by default 'Time Series Data'.
    ylabel : str, optional
        Label for the y-axis, by default 'Value'.

    Returns
    -------
    dict
        The updated state dictionary.
    """
    logger.info(f"Plotting time series data: {title}...")
    try:
        df = state['data']['original']
        if df is None:
            logger.warning("No original time series data found to plot.")
            state['logs'].append("WARNING: No original time series data found for plotting.")
            return state

        plt.figure(figsize=(12, 6))
        plt.plot(df.index, df['value'], label='Original Series')
        plt.title(title)
        plt.xlabel('Date')
        plt.ylabel(ylabel)
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.show()
        logger.info("Time series plot displayed.")
        state['logs'].append("Time series data plotted.")
    except Exception as e:
        logger.error(f"Error plotting time series data: {e}")
        state['logs'].append(f"ERROR: Failed to plot time series data - {e}")
    return state

def split_data(state: dict) -> dict:
    """
    Splits the original time series data into training and testing sets.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting 'original' time series data
        and 'test_size' in config.

    Returns
    -------
    dict
        The updated state dictionary with 'train' and 'test' dataframes.
    """
    logger.info("Splitting data into training and testing sets...")
    try:
        original_data = state['data']['original']
        if original_data is None:
            logger.error("No original data found for splitting. Please generate data first.")
            state['logs'].append("ERROR: No original data to split.")
            return state

        test_size = state['config']['test_size']
        train_size = int(len(original_data) * (1 - test_size))

        train_data = original_data.iloc[:train_size]
        test_data = original_data.iloc[train_size:]

        state['data']['train'] = train_data
        state['data']['test'] = test_data

        logger.info(f"Data split: Training set size = {len(train_data)}, Test set size = {len(test_data)}.")
        state['logs'].append(f"Data split with {len(train_data)} training and {len(test_data)} test samples.")
    except Exception as e:
        logger.error(f"Error splitting data: {e}")
        state['logs'].append(f"ERROR: Failed to split data - {e}")
    return state

def check_stationarity(state: dict, series_name: str = 'Time Series Data') -> dict:
    """
    Checks the stationarity of a time series using the Augmented Dickey-Fuller test
    and visualizes ACF/PACF plots.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting the time series data under state['data']['train']
        or a specific key for differenced data.
    series_name : str, optional
        A descriptive name for the series being checked, by default 'Time Series Data'.

    Returns
    -------
    dict
        The updated state dictionary with stationarity test results logged.
    """
    logger.info(f"Checking stationarity for {series_name}...")
    try:
        data = state['data']['train']['value'] # Default to training data for initial check

        # Perform Augmented Dickey-Fuller test
        adf_test = adfuller(data.dropna())
        adf_output = pd.Series(adf_test[0:4], index=['Test Statistic', 'p-value', '#Lags Used', 'Number of Observations Used'])
        for key, value in adf_test[4].items():
            adf_output[f'Critical Value ({key})'] = value

        logger.info(f"ADF Test Results for {series_name}:\n{adf_output}")
        state['logs'].append(f"ADF Test for {series_name}: p-value = {adf_output['p-value']:.4f}")

        if adf_output['p-value'] <= 0.05:
            logger.info(f"The {series_name} is likely stationary (p-value <= 0.05).")
            state['logs'].append(f"INFO: {series_name} is stationary.")
        else:
            logger.warning(f"The {series_name} is likely non-stationary (p-value > 0.05). Differencing may be required.")
            state['logs'].append(f"WARNING: {series_name} is non-stationary.")

        # Plot time series, ACF and PACF
        fig, axes = plt.subplots(3, 1, figsize=(14, 12))

        # Plot original series
        axes[0].plot(data)
        axes[0].set_title(f'{series_name} Time Series')
        axes[0].set_ylabel('Value')
        axes[0].grid(True)

        # Plot ACF
        plot_acf(data.dropna(), ax=axes[1], lags=min(len(data)//2 -1, 40))
        axes[1].set_title(f'{series_name} Autocorrelation Function (ACF)')
        axes[1].grid(True)

        # Plot PACF
        plot_pacf(data.dropna(), ax=axes[2], lags=min(len(data)//2 -1, 40))
        axes[2].set_title(f'{series_name} Partial Autocorrelation Function (PACF)')
        axes[2].grid(True)

        plt.tight_layout()
        plt.show()

        logger.info(f"Stationarity plots for {series_name} displayed.")
        state['logs'].append(f"Stationarity plots for {series_name} generated.")

    except Exception as e:
        logger.error(f"Error checking stationarity for {series_name}: {e}")
        state['logs'].append(f"ERROR: Failed to check stationarity for {series_name} - {e}")
    return state

def apply_differencing(state: dict, order: int = 1) -> dict:
    """
    Applies differencing to the time series data to achieve stationarity.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting 'train' data.
    order : int, optional
        The order of differencing to apply, by default 1.

    Returns
    -------
    dict
        The updated state dictionary with the differenced data.
    """
    logger.info(f"Applying differencing of order {order} to the training data...")
    try:
        train_data = state['data']['train']['value']
        if train_data is None:
            logger.error("No training data found for differencing.")
            state['logs'].append("ERROR: No training data for differencing.")
            return state

        differenced_data = train_data.diff(periods=order).dropna()
        state['data']['stationary'] = differenced_data
        logger.info(f"Data differenced {order} time(s). First 5 values: {differenced_data.head().tolist()}")
        state['logs'].append(f"Differencing of order {order} applied.")

        # Re-check stationarity after differencing
        state = check_stationarity(state, series_name=f'Differenced Data (Order {order})')

    except Exception as e:
        logger.error(f"Error applying differencing: {e}")
        state['logs'].append(f"ERROR: Failed to apply differencing - {e}")
    return state

def fit_arima_model(state: dict, order: tuple[int, int, int]) -> dict:
    """
    Fits an ARIMA model to the training data with specified (p, d, q) order.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting 'train' data and configuration for retries.
    order : tuple[int, int, int]
        The (p, d, q) order for the ARIMA model.

    Returns
    -------
    dict
        The updated state dictionary with the fitted ARIMA model.
    """
    p, d, q = order
    logger.info(f"Attempting to fit ARIMA model with order ({p}, {d}, {q})...")
    train_data = state['data']['train']['value']

    if train_data is None:
        logger.error("No training data available to fit the ARIMA model.")
        state['logs'].append("ERROR: No training data for ARIMA model fitting.")
        return state

    retries = state['config'].get('retries', 3)
    backoff_factor = state['config'].get('backoff_factor', 0.5)

    for i in range(retries):
        try:
            model = ARIMA(train_data, order=order)
            fitted_model = model.fit()
            state['model']['fitted_model'] = fitted_model
            state['model']['arima_order'] = order
            logger.info(f"ARIMA model with order ({p}, {d}, {q}) fitted successfully.")
            state['logs'].append(f"ARIMA model fitted with order {order}.")
            logger.debug(f"ARIMA model summary: {fitted_model.summary()}")
            return state
        except Exception as e:
            delay = (backoff_factor * (2 ** i)) + (random.random() * 0.1)
            logger.warning(f"Attempt {i+1}/{retries} failed to fit ARIMA model with order {order}: {e}. Retrying in {delay:.2f} seconds...")
            state['logs'].append(f"WARNING: ARIMA fit attempt {i+1} failed: {e}.")
            time.sleep(delay)

    logger.error(f"Failed to fit ARIMA model with order {order} after {retries} attempts.")
    state['logs'].append(f"ERROR: Failed to fit ARIMA model after {retries} attempts for order {order}.")
    return state

def make_predictions(state: dict, forecast_steps: int = None) -> dict:
    """
    Generates predictions using the fitted ARIMA model.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting a fitted ARIMA model.
    forecast_steps : int, optional
        The number of steps to forecast. If None, it defaults to the length of the test data.

    Returns
    -------
    dict
        The updated state dictionary with predictions.
    """
    logger.info("Making predictions with the fitted ARIMA model...")
    fitted_model = state['model'].get('fitted_model')
    test_data = state['data'].get('test')

    if fitted_model is None:
        logger.error("No fitted ARIMA model found to make predictions.")
        state['logs'].append("ERROR: No fitted model for predictions.")
        return state

    if test_data is None:
        logger.warning("No test data available. Predictions will be generated based on forecast_steps.")
        # If no test data, ensure forecast_steps is provided
        if forecast_steps is None:
            logger.error("Cannot make predictions without test data or specified forecast_steps.")
            state['logs'].append("ERROR: No test data or forecast_steps for predictions.")
            return state

    if forecast_steps is None:
        forecast_steps = len(test_data)

    try:
        # Start prediction from the end of the training data
        start_idx = len(state['data']['train'])
        end_idx = start_idx + forecast_steps - 1

        # Get the index for the forecast period using the original data's index
        # This handles cases where original data might have gaps or different frequencies than simple integer range
        forecast_index = state['data']['original'].index[start_idx : end_idx + 1]

        if len(forecast_index) != forecast_steps:
             logger.warning(f"Forecast index length ({len(forecast_index)}) does not match requested steps ({forecast_steps}). Adjusting forecast_steps.")
             forecast_steps = len(forecast_index)
             end_idx = start_idx + forecast_steps - 1

        predictions_series = fitted_model.predict(start=start_idx, end=end_idx, typ='levels')

        # Ensure the predictions have the correct index
        predictions_series.index = forecast_index

        state['results']['predictions'] = predictions_series
        state['model']['forecast_steps'] = forecast_steps
        logger.info(f"Predictions made for {forecast_steps} steps.")
        state['logs'].append(f"Predictions generated for {forecast_steps} steps.")

    except Exception as e:
        logger.error(f"Error making predictions: {e}")
        state['logs'].append(f"ERROR: Failed to make predictions - {e}")
    return state

def evaluate_model(state: dict) -> dict:
    """
    Evaluates the performance of the ARIMA model's predictions.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting predictions and test data.

    Returns
    -------
    dict
        The updated state dictionary with evaluation metrics.
    """
    logger.info("Evaluating model predictions...")
    predictions = state['results'].get('predictions')
    test_data = state['data'].get('test')

    if predictions is None:
        logger.error("No predictions found to evaluate.")
        state['logs'].append("ERROR: No predictions for evaluation.")
        return state

    if test_data is None:
        logger.error("No test data found for evaluation.")
        state['logs'].append("ERROR: No test data for evaluation.")
        return state

    try:
        # Align predictions and test data by index to ensure proper comparison
        # This handles cases where prediction length might differ slightly due to index issues
        common_index = predictions.index.intersection(test_data.index)
        aligned_predictions = predictions.loc[common_index]
        aligned_test_data = test_data.loc[common_index, 'value']

        if len(common_index) == 0:
            logger.error("No overlapping indices between predictions and test data. Cannot evaluate.")
            state['logs'].append("ERROR: No overlapping data for evaluation.")
            return state

        mse = mean_squared_error(aligned_test_data, aligned_predictions)
        rmse = np.sqrt(mse)
        mae = np.mean(np.abs(aligned_test_data - aligned_predictions))

        state['results']['metrics']['mse'] = mse
        state['results']['metrics']['rmse'] = rmse
        state['results']['metrics']['mae'] = mae

        logger.info(f"Model Evaluation Results: MSE={mse:.4f}, RMSE={rmse:.4f}, MAE={mae:.4f}")
        state['logs'].append(f"Model evaluated: RMSE={rmse:.4f}")

    except Exception as e:
        logger.error(f"Error evaluating model: {e}")
        state['logs'].append(f"ERROR: Failed to evaluate model - {e}")
    return state

def plot_predictions(state: dict, title: str = 'ARIMA Model Predictions') -> dict:
    """
    Plots the historical data, test data, and model predictions.

    Parameters
    ----------
    state : dict
        The current state dictionary, expecting original, train, test data and predictions.
    title : str, optional
        Title of the plot, by default 'ARIMA Model Predictions'.

    Returns
    -------
    dict
        The updated state dictionary.
    """
    logger.info("Plotting predictions against actuals...")
    original_data = state['data'].get('original')
    train_data = state['data'].get('train')
    test_data = state['data'].get('test')
    predictions = state['results'].get('predictions')

    if original_data is None or train_data is None or test_data is None or predictions is None:
        logger.error("Missing required data (original, train, test, or predictions) for plotting.")
        state['logs'].append("ERROR: Missing data for predictions plot.")
        return state

    try:
        plt.figure(figsize=(15, 7))
        plt.plot(train_data.index, train_data['value'], label='Training Data', color='blue')
        plt.plot(test_data.index, test_data['value'], label='Actual Test Data', color='green')
        plt.plot(predictions.index, predictions, label='ARIMA Predictions', color='red', linestyle='--')

        plt.title(title)
        plt.xlabel('Date')
        plt.ylabel('Value')
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.show()
        logger.info("Prediction plot displayed.")
        state['logs'].append("Prediction plot generated.")

    except Exception as e:
        logger.error(f"Error plotting predictions: {e}")
        state['logs'].append(f"ERROR: Failed to plot predictions - {e}")
    return state

def run_arima_workflow(
    start_date: str = '2020-01-01',
    end_date: str = '2023-01-01',
    frequency: str = 'D',
    ar_params: list[float] = None,
    ma_params: list[float] = None,
    n_samples: int = 1000,
    arima_order: tuple[int, int, int] = (1, 1, 1),
    test_size: float = 0.2,
    random_seed: int = 42
) -> dict:
    """
    Orchestrates the entire ARIMA time series forecasting workflow.

    Parameters
    ----------
    start_date : str, optional
        Start date for data generation, by default '2020-01-01'.
    end_date : str, optional
        End date for data generation, by default '2023-01-01'.
    frequency : str, optional
        Frequency of time series data, by default 'D'.
    ar_params : list[float], optional
        AR coefficients for data generation, by default None.
    ma_params : list[float], optional
        MA coefficients for data generation, by default None.
    n_samples : int, optional
        Number of samples for data generation, by default 1000.
    arima_order : tuple[int, int, int], optional
        The (p, d, q) order for the ARIMA model, by default (1, 1, 1).
    test_size : float, optional
        Proportion of data to reserve for testing, by default 0.2.
    random_seed : int, optional
        Seed for reproducibility, by default 42.

    Returns
    -------
    dict
        The final state dictionary containing all results and configurations.
    """
    logger.info("Starting ARIMA workflow...")
    current_state = create_state(
        start_date=start_date, end_date=end_date, frequency=frequency,
        ar_params=ar_params, ma_params=ma_params, n_samples=n_samples,
        random_seed=random_seed
    )
    current_state['config']['test_size'] = test_size # Override default test_size if provided

    current_state = create_time_series_data(current_state)
    current_state = plot_time_series(current_state, title='Generated Time Series Data')
    current_state = split_data(current_state)
    current_state = check_stationarity(current_state, series_name='Training Data')

    # Apply differencing based on the 'd' order from arima_order
    d_order = arima_order[1]
    if d_order > 0:
        logger.info(f"Applying differencing of order {d_order} as per ARIMA order...")
        current_state = apply_differencing(current_state, order=d_order)
    else:
        logger.info("No differencing applied (d=0).")

    current_state = fit_arima_model(current_state, order=arima_order)

    # Ensure model was fitted before proceeding with predictions and evaluation
    if current_state['model']['fitted_model'] is not None:
        current_state = make_predictions(current_state)
        current_state = evaluate_model(current_state)
        current_state = plot_predictions(current_state)
    else:
        logger.error("ARIMA model was not fitted successfully. Skipping predictions and evaluation.")
        current_state['logs'].append("ERROR: ARIMA workflow aborted due to model fit failure.")

    logger.info("ARIMA workflow completed.")
    for log_entry in current_state['logs']:
        logger.info(f"Workflow Log: {log_entry}")
    return current_state

# Example usage of the ARIMA workflow
# Define parameters for the workflow
workflow_params = {
    'start_date': '2015-01-01',
    'end_date': '2023-01-01',
    'frequency': 'D',
    'ar_params': [0.8, -0.2],
    'ma_params': [0.5, 0.3],
    'n_samples': 2000,
    'arima_order': (2, 1, 2), # (p, d, q) example
    'test_size': 0.15,
    'random_seed': 42
}

# Run the complete ARIMA workflow
final_state = run_arima_workflow(**workflow_params)

print("\n--- Final State Summary ---")
print(f"ARIMA Order Used: {final_state['model']['arima_order']}")
if final_state['results']['metrics']:
    print(f"RMSE: {final_state['results']['metrics'].get('rmse', 'N/A'):.4f}")
    print(f"MAE: {final_state['results']['metrics'].get('mae', 'N/A'):.4f}")
else:
    print("Model evaluation metrics not available.")

print("\n--- Workflow Logs ---")
for log in final_state['logs']:
    print(log)