# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
# See the LICENSE file in the root of this repository for details.
# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial

from dataclasses import dataclass

import pandas as pd
from arch.unitroot import PhillipsPerron
from numpy.linalg import LinAlgError

from validmind.logging import get_logger
from validmind.vm_models import Metric, ResultSummary, ResultTable, ResultTableMetadata

logger = get_logger(__name__)


@dataclass
class PhillipsPerronArch(Metric):
    """
    Assesses the stationarity of time series data in each feature of the ML model using the Phillips-Perron test.

    ### Purpose

    The Phillips-Perron (PP) test is used to determine the stationarity of time series data for each feature in a
    dataset, which is crucial for forecasting tasks. It tests the null hypothesis that a time series is unit-root
    non-stationary. This is vital for understanding the stochastic behavior of the data and ensuring the robustness and
    validity of predictions generated by regression analysis models.

    ### Test Mechanism

    The PP test is conducted for each feature in the dataset as follows:
    - A data frame is created from the dataset.
    - For each column, the Phillips-Perron method calculates the test statistic, p-value, lags used, and number of
    observations.
    - The results are then stored for each feature, providing a metric that indicates the stationarity of the time
    series data.

    ### Signs of High Risk

    - A high p-value, indicating that the series has a unit root and is non-stationary.
    - Test statistic values exceeding critical values, suggesting non-stationarity.
    - High 'usedlag' value, pointing towards autocorrelation issues that may degrade model performance.

    ### Strengths

    - Resilience against heteroskedasticity in the error term.
    - Effective for long time series data.
    - Helps in determining whether the time series is stationary, aiding in the selection of suitable forecasting
    models.

    ### Limitations

    - Applicable only within a univariate time series framework.
    - Relies on asymptotic theory, which may reduce the test’s power for small sample sizes.
    - Non-stationary time series must be converted to stationary series through differencing, potentially leading to
    loss of important data points.
    """

    name = "phillips_perron"
    required_inputs = ["dataset"]
    tasks = ["regression"]
    tags = [
        "time_series_data",
        "forecasting",
        "statistical_test",
        "unit_root_test",
    ]

    def run(self):
        """
        Calculates PP metric for each of the dataset features
        """
        dataset = self.inputs.dataset.df

        # Check if the dataset is a time series
        if not isinstance(dataset.index, (pd.DatetimeIndex, pd.PeriodIndex)):
            raise ValueError(
                "Dataset index must be a datetime or period index for time series analysis."
            )

        # Preprocessing: Drop rows with any NaN values
        if dataset.isnull().values.any():
            logger.warning(
                "Dataset contains missing values. Rows with NaNs will be dropped."
            )
            dataset = dataset.dropna()

        # Convert to numeric and handle non-numeric data
        dataset = dataset.apply(pd.to_numeric, errors="coerce")

        # Initialize a list to store Phillips-Perron results
        pp_values = []

        for col in dataset.columns:
            try:
                pp = PhillipsPerron(dataset[col].values)
                pp_values.append(
                    {
                        "Variable": col,
                        "stat": pp.stat,
                        "pvalue": pp.pvalue,
                        "usedlag": pp.lags,
                        "nobs": pp.nobs,
                    }
                )
            except LinAlgError as e:
                logger.error(f"Error processing column '{col}': {e}")
                pp_values.append(
                    {
                        "Variable": col,
                        "stat": None,
                        "pvalue": None,
                        "usedlag": None,
                        "nobs": None,
                        "error": str(e),
                    }
                )

        return self.cache_results({"phillips_perron_results": pp_values})

    def summary(self, metric_value):
        """
        Build a table for summarizing the Phillips-Perron results
        """
        pp_results = metric_value["phillips_perron_results"]

        return ResultSummary(
            results=[
                ResultTable(
                    data=pp_results,
                    metadata=ResultTableMetadata(title="Phillips-Perron Test Results"),
                )
            ]
        )
