path: root/openbb_platform/extensions/quantitative/openbb_quantitative/quantitative_router.py

"""Quantitative Analysis Router."""

# pylint: disable=unused-argument

from typing import List, Literal

import pandas as pd
from openbb_core.app.model.example import APIEx, PythonEx
from openbb_core.app.model.obbject import OBBject
from openbb_core.app.router import Router
from openbb_core.app.utils import (
    basemodel_to_df,
    get_target_column,
    get_target_columns,
)
from openbb_core.provider.abstract.data import Data

from openbb_quantitative.performance.performance_router import (
    router as performance_router,
)
from openbb_quantitative.rolling.rolling_router import router as rolling_router
from openbb_quantitative.stats.stats_router import router as stats_router

from .helpers import get_fama_raw
from .models import (
    ADFTestModel,
    CAPMModel,
    KPSSTestModel,
    NormalityModel,
    SummaryModel,
    TestModel,
    UnitRootModel,
)

router = Router(prefix="")
router.include_router(rolling_router)
router.include_router(stats_router)
router.include_router(performance_router)


@router.command(
    methods=["POST"],
    examples=[
        PythonEx(
            description="Get Normality Statistics.",
            code=[
                "stock_data = obb.equity.price.historical(symbol='TSLA', start_date='2023-01-01', provider='fmp').to_df()",  # noqa: E501
                "obb.quantitative.normality(data=stock_data, target='close')",
            ],
        ),
        APIEx(parameters={"target": "close", "data": APIEx.mock_data("timeseries", 8)}),
    ],
)
def normality(
    data: List[Data],
    target: str,
    **extra_params,
) -> OBBject[NormalityModel]:
    """Get Normality Statistics.

    - **Kurtosis**: whether the kurtosis of a sample differs from the normal distribution.
    - **Skewness**: whether the skewness of a sample differs from the normal distribution.
    - **Jarque-Bera**: whether the sample data has the skewness and kurtosis matching a normal distribution.
    - **Shapiro-Wilk**: whether a random sample comes from a normal distribution.
    - **Kolmogorov-Smirnov**: whether two underlying one-dimensional probability distributions differ.

    Parameters
    ----------
    data : List[Data]
        Time series data.
    target : str
        Target column name.
    **extra_params : Optional[Dict[str, Any]]
        Extra parameters to be passed to the command execution.
        API POST requests are sent in the body with data.

    Returns
    -------
    OBBject[NormalityModel]
        Normality tests summary. See qa_models.NormalityModel for details.
    """
    from scipy import stats  # pylint: disable=import-outside-toplevel

    df = basemodel_to_df(data)
    series_target = get_target_column(df, target)

    kt_statistic, kt_pvalue = stats.kurtosistest(series_target)
    sk_statistic, sk_pvalue = stats.skewtest(series_target)
    jb_statistic, jb_pvalue = stats.jarque_bera(series_target)
    sh_statistic, sh_pvalue = stats.shapiro(series_target)
    ks_statistic, ks_pvalue = stats.kstest(series_target, "norm")

    norm_summary = NormalityModel(
        kurtosis=TestModel(statistic=kt_statistic, p_value=kt_pvalue),
        skewness=TestModel(statistic=sk_statistic, p_value=sk_pvalue),
        jarque_bera=TestModel(statistic=jb_statistic, p_value=jb_pvalue),
        shapiro_wilk=TestModel(statistic=sh_statistic, p_value=sh_pvalue),
        kolmogorov_smirnov=TestModel(statistic=ks_statistic, p_value=ks_pvalue),
    )

    return OBBject(results=norm_summary)


@router.command(
    methods=["POST"],
    examples=[
        PythonEx(
            description="Get Capital Asset Pricing Model (CAPM).",
            code=[
                "stock_data = obb.equity.price.historical(symbol='TSLA', start_date='2023-01-01', provider='fmp').to_df()",  # noqa: E501
                "obb.quantitative.capm(data=stock_data, target='close')",
            ],
        ),
        APIEx(
            parameters={"target": "close", "data": APIEx.mock_data("timeseries", 31)}
        ),
    ],
)
def capm(
    data: List[Data],
    target: str,
    **extra_params,
) -> OBBject[CAPMModel]:
    """Get Capital Asset Pricing Model (CAPM).

    CAPM offers a streamlined way to assess the expected return on an investment while accounting for its risk relative
    to the market. It's a cornerstone of modern financial theory that helps investors understand the trade-off between
    risk and return, guiding more informed investment choices.

    Parameters
    ----------
    data : List[Data]
        Time series data.
    target : str
        Target column name.
    **extra_params : Optional[Dict[str, Any]]
        Extra parameters to be passed to the command execution.
        API POST requests are sent in the body with data.

    Returns
    -------
    OBBject[CAPMModel]
        CAPM model summary.
    """
    import statsmodels.api as sm  # pylint: disable=import-outside-toplevel # type: ignore

    df = basemodel_to_df(data)

    df_target = get_target_columns(df, ["date", target])
    df_target = df_target.set_index("date")
    df_target.loc[:, "return"] = df_target.pct_change()
    df_target = df_target.dropna()
    df_target.index = pd.to_datetime(df_target.index)
    start_date = df_target.index.min().strftime("%Y-%m-%d")
    end_date = df_target.index.max().strftime("%Y-%m-%d")
    df_fama = get_fama_raw(start_date, end_date)
    df_target = df_target.merge(df_fama, left_index=True, right_index=True)
    df_target["excess_return"] = df_target["return"] - df_target["RF"]
    df_target["excess_mkt"] = df_target["MKT-RF"] - df_target["RF"]
    df_target = df_target.dropna()

    y = df_target[["excess_return"]]
    x = df_target["excess_mkt"]
    x = sm.add_constant(x)
    model = sm.OLS(y, x).fit()

    results = CAPMModel(
        market_risk=model.params["excess_mkt"],
        systematic_risk=model.rsquared,
        idiosyncratic_risk=1 - model.rsquared,
    )

    return OBBject(results=results)


@router.command(
    methods=["POST"],
    examples=[
        PythonEx(
            description="Get Unit Root Test.",
            code=[
                "stock_data = obb.equity.price.historical(symbol='TSLA', start_date='2023-01-01', provider='fmp').to_df()",  # noqa: E501
                "obb.quantitative.unitroot_test(data=stock_data, target='close')",
            ],
        ),
        APIEx(parameters={"target": "close", "data": APIEx.mock_data("timeseries", 5)}),
    ],
)
def unitroot_test(
    data: List[Data],
    target: str,
    fuller_reg: Literal["c", "ct", "ctt", "nc", "c"] = "c",
    kpss_reg: Literal["c", "ct"] = "c",
    **extra_params,
) -> OBBject[UnitRootModel]:
    """Get Unit Root Test.

    This function applies two renowned tests to assess whether your data series is stationary or if it contains a unit
    root, indicating it may be influenced by time-based trends or seasonality. The Augmented Dickey-Fuller (ADF) test
    helps identify the presence of a unit root, suggesting that the series could be non-stationary and potentially
    unpredictable over time. On the other hand, the Kwiatkowski-Phillips-Schmidt-Shin (KPSS) test checks for the
    stationarity of the series, where failing to reject the null hypothesis indicates a stable, stationary series.
    Together, these tests provide a comprehensive view of your data's time series properties, essential for
    accurate modeling and forecasting.

    Parameters
    ----------
    data : List[Data]
        Time series data.
    target : str
        Target column name.
    fuller_reg : Literal["c", "ct", "ctt", "nc", "c"]
        Regression type for ADF test.
    kpss_reg : Literal["c", "ct"]
        Regression type for KPSS test.
    **extra_params : Optional[Dict[str, Any]]
        Extra parameters to be passed to the command execution.
        API POST requests are sent in the body with data.

    Returns
    -------
    OBBject[UnitRootModel]
        Unit root tests summary.
    """
    # pylint: disable=import-outside-toplevel
    from statsmodels.tsa import stattools  # type: ignore

    df = basemodel_to_df(data)
    series_target = get_target_column(df, target)

    adf = stattools.adfuller(series_target, regression=fuller_reg)
    kpss = stattools.kpss(series_target, regression=kpss_reg, nlags="auto")

    unitroot_summary = UnitRootModel(
        adf=ADFTestModel(
            statistic=adf[0],
            p_value=adf[1],
            nlags=adf[2] if isinstance(adf[2], int) else 0,
            nobs=adf[3] if isinstance(adf[3], int) else 0,
            icbest=adf[5] if isinstance(adf[5], float) else 0.0,  # type: ignore
        ),
        kpss=KPSSTestModel(
            statistic=kpss[0],
            p_value=kpss[1],
            nlags=kpss[2],
        ),
    )
    return OBBject(results=unitroot_summary)


@router.command(
    methods=["POST"],
    examples=[
        PythonEx(
            description="Get Summary Statistics.",
            code=[
                "stock_data = obb.equity.price.historical(symbol='TSLA', start_date='2023-01-01', provider='fmp').to_df()",  # noqa: E501
                "obb.quantitative.summary(data=stock_data, target='close')",
            ],
        ),
        APIEx(parameters={"target": "close", "data": APIEx.mock_data("timeseries", 5)}),
    ],
)
def summary(
    data: List[Data],
    target: str,
    **extra_params,
) -> OBBject[SummaryModel]:
    """Get Summary Statistics.

    The summary that offers a snapshot of its central tendencies, variability, and distribution.
    This command calculates essential statistics, including mean, standard deviation, variance,
    and specific percentiles, to provide a detailed profile of your target column. B
    y examining these metrics, you gain insights into the data's overall behavior, helping to identify patterns,
    outliers, or anomalies. The summary table is an invalua