OpenBB/openbb_terminal/common/prediction_techniques/arima_model.py

"""Arima Prediction Model"""
__docformat__ = "numpy"

import logging
from typing import Any, List, Tuple, Union

import pandas as pd
import pmdarima
from statsmodels.tsa.arima.model import ARIMA

from openbb_terminal.decorators import log_start_end

logger = logging.getLogger(__name__)

ICS = ["aic", "aicc", "bic", "hqic", "oob"]


@log_start_end(log=logger)
def get_arima_model(
    values: Union[pd.Series, pd.DataFrame],
    arima_order: str,
    n_predict: int,
    seasonal: bool,
    ic: str,
) -> Tuple[List[float], Any]:
    """Get an ARIMA model for data

    Parameters
    ----------
    values : Union[pd.Series, pd.DataFrame]
        Data to fit
    arima_order : str
        String of ARIMA params in form "p,q,d"
    n_predict : int
        Days to predict
    seasonal : bool
        Flag to use seasonal model
    ic : str
        Information Criteria for model evaluation

    Returns
    -------
    List[float]
        List of predicted values
    Any
        Fit ARIMA model object.
    """
    if arima_order:
        model = ARIMA(
            values, order=tuple(int(ord) for ord in arima_order.split(","))
        ).fit()
        l_predictions = list(
            model.predict(
                start=len(values.values) + 1,
                end=len(values.values) + n_predict,
            )
        )
    else:
        if seasonal:
            model = pmdarima.auto_arima(
                values.values,
                error_action="ignore",
                seasonal=True,
                m=5,
                information_criteria=ic,
            )
        else:
            model = pmdarima.auto_arima(
                values.values,
                error_action="ignore",
                seasonal=False,
                information_criteria=ic,
            )
        l_predictions = list(model.predict(n_predict))

    return l_predictions, model