"""
Frees test for cross-sectional dependence in panel data.

References
----------
Frees, E. W. (1995). Assessing Cross-Sectional Correlation in Panel Data.
Journal of Econometrics, 69(2), 393-414.

Frees, E. W. (2004). Longitudinal and Panel Data: Analysis and Applications
in the Social Sciences. Cambridge University Press.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Union

if TYPE_CHECKING:
    from panelbox.core.results import PanelResults

import numpy as np
import pandas as pd
from scipy import stats

from panelbox.validation.base import ValidationTest, ValidationTestResult


class FreesTest(ValidationTest):
    """
    Frees test for cross-sectional dependence in panel data.

    This is a non-parametric test that uses the Friedman-type rank
    statistic to test for cross-sectional dependence.

    H0: No cross-sectional dependence (residuals independent across entities)
    H1: Cross-sectional dependence present

    The test is robust to non-normality and heteroskedasticity.

    Notes
    -----
    The test statistic is based on the average Spearman rank correlation:

    Q_F = sum_{i<j} R_ij² / [N(N-1)/2]

    where R_ij is the Spearman rank correlation between residuals
    of entities i and j.

    Under H0, Q_F follows an approximate distribution that can be
    compared to critical values tabulated by Frees (1995, 2004).

    For large samples, we use asymptotic approximation:
    sqrt(T) * Q_F is approximately normal.

    Advantages over Breusch-Pagan LM test:
    - Non-parametric (doesn't require normality)
    - More robust to outliers
    - Better size properties for moderate N

    Examples
    --------
    >>> from panelbox.models.static.pooled_ols import PooledOLS
    >>> model = PooledOLS("y ~ x1 + x2", data, "entity", "time")
    >>> results = model.fit()
    >>>
    >>> from panelbox.validation.cross_sectional_dependence.frees import FreesTest
    >>> test = FreesTest(results)
    >>> result = test.run()
    >>> print(result)
    """

    def __init__(self, results: "PanelResults"):
        """
        Initialize Frees test.

        Parameters
        ----------
        results : PanelResults
            Results from panel model estimation
        """
        super().__init__(results)

    def run(self, alpha: float = 0.05, **kwargs) -> ValidationTestResult:
        """
        Run Frees test for cross-sectional dependence.

        Parameters
        ----------
        alpha : float, default=0.05
            Significance level

        Returns
        -------
        ValidationTestResult
            Test results

        Notes
        -----
        The test uses asymptotic approximation for p-values.
        Critical values from Frees (1995) are also provided in metadata
        for reference.
        """
        # Get residuals with entity and time structure
        resid_df = self._prepare_residual_data()

        # Create wide format: rows = time, columns = entities
        resid_wide = resid_df.pivot(index="time", columns="entity", values="resid")

        # Get dimensions
        T = resid_wide.shape[0]  # Number of time periods
        N = resid_wide.shape[1]  # Number of entities

        if N < 2:
            raise ValueError("Need at least 2 entities for cross-sectional dependence test")

        if T < 3:
            raise ValueError("Need at least 3 time periods for Frees test")

        # Compute pairwise Spearman correlations
        rank_correlations: list[float] = []
        n_pairs = 0

        entity_list = list(resid_wide.columns)

        for i in range(N):
            for j in range(i + 1, N):
                entity_i = entity_list[i]
                entity_j = entity_list[j]

                # Get residuals for this pair (drop NaN)
                resid_i = resid_wide[entity_i].dropna()
                resid_j = resid_wide[entity_j].dropna()

                # Find common time periods
                common_times = resid_i.index.intersection(resid_j.index)

                if len(common_times) >= 3:  # Need at least 3 obs for rank correlation
                    resid_i_common = resid_i.loc[common_times].values
                    resid_j_common = resid_j.loc[common_times].values

                    # Compute Spearman rank correlation
                    # This is robust to outliers and non-normality
                    rho_rank, _ = stats.spearmanr(resid_i_common, resid_j_common)

                    # Handle potential NaN from constant series
                    if not np.isnan(rho_rank):
                        rank_correlations.append(rho_rank)
                        n_pairs += 1

        if n_pairs == 0:
            raise ValueError(
                "No valid pairwise rank correlations could be computed. "
                "Check for constant residuals or insufficient data."
            )

        # Compute Frees statistic
        # Q_F = mean of squared rank correlations
        rank_correlations = np.array(rank_correlations)
        q_frees = np.mean(rank_correlations**2)

        # Asymptotic distribution
        # Under H0: E[Q_F] = 1/(T-1)
        # Var[Q_F] ≈ 2(T-3) / [(T+1)(T-1)²]
        expected_qf = 1 / (T - 1)
        var_qf = 2 * (T - 3) / ((T + 1) * (T - 1) ** 2)
        se_qf = np.sqrt(var_qf / n_pairs)  # SE of mean

        # Test statistic: standardized Q_F
        if se_qf > 0:
            z_stat = (q_frees - expected_qf) / se_qf
        else:
            z_stat = 0.0

        # P-value (two-sided)
        pvalue = 2 * (1 - stats.norm.cdf(abs(z_stat)))

        # Critical values from Frees (1995) Table 1
        # These are for balanced panels at alpha=0.05
        critical_values = self._get_critical_values(T, N)

        # Simple interpretation
        interpretation = (
            "Reject H0 (cross-sectional dependence detected)"
            if q_frees > critical_values.get("alpha_0.05", float("inf"))
            else "Do not reject H0 (no evidence of cross-sectional dependence)"
        )

        # Metadata
        mean_abs_rank_corr = np.mean(np.abs(rank_correlations))
        max_abs_rank_corr = np.max(np.abs(rank_correlations))

        metadata = {
            "q_frees_statistic": float(q_frees),
            "z_statistic": float(z_stat),
            "expected_qf_under_h0": float(expected_qf),
            "n_entities": int(N),
            "n_time_periods": int(T),
            "n_pairs": int(n_pairs),
            "mean_abs_rank_correlation": float(mean_abs_rank_corr),
            "max_abs_rank_correlation": float(max_abs_rank_corr),
            "critical_values": critical_values,
            "interpretation": interpretation,
            "note": (
                "Frees test is non-parametric and robust to non-normality. "
                "Critical values are approximate for unbalanced panels."
            ),
        }

        result = ValidationTestResult(
            test_name="Frees Test for Cross-Sectional Dependence",
            statistic=z_stat,
            pvalue=pvalue,
            null_hypothesis="No cross-sectional dependence (residuals independent across entities)",
            alternative_hypothesis="Cross-sectional dependence present",
            alpha=alpha,
            df=None,  # Non-parametric test
            metadata=metadata,
        )

        return result

    def _get_critical_values(self, T, N):
        """
        Get critical values from Frees (1995) Table 1.

        Parameters
        ----------
        T : int
            Number of time periods
        N : int
            Number of entities

        Returns
        -------
        dict
            Critical values at different significance levels

        Notes
        -----
        These are approximate critical values for balanced panels.
        Interpolation is used for T values not in the table.
        """
        # Approximate critical values from Frees (1995, 2004)
        # Q_F critical values at alpha = 0.10, 0.05, 0.01
        # Rows: T, Columns: significance levels

        # For small T, use tabulated values
        if T <= 5:
            cv_0_10 = 0.3000  # Very approximate
            cv_0_05 = 0.4000
            cv_0_01 = 0.6000
        elif T <= 10:
            cv_0_10 = 0.1429
            cv_0_05 = 0.2000
            cv_0_01 = 0.3000
        elif T <= 20:
            cv_0_10 = 0.0800
            cv_0_05 = 0.1100
            cv_0_01 = 0.1700
        elif T <= 30:
            cv_0_10 = 0.0543
            cv_0_05 = 0.0754
            cv_0_01 = 0.1170
        else:
            # Asymptotic approximation
            cv_0_10 = 1.28 * np.sqrt(2 * (T - 3) / ((T + 1) * (T - 1) ** 2)) + 1 / (T - 1)
            cv_0_05 = 1.96 * np.sqrt(2 * (T - 3) / ((T + 1) * (T - 1) ** 2)) + 1 / (T - 1)
            cv_0_01 = 2.58 * np.sqrt(2 * (T - 3) / ((T + 1) * (T - 1) ** 2)) + 1 / (T - 1)

        return {
            "alpha_0.10": cv_0_10,
            "alpha_0.05": cv_0_05,
            "alpha_0.01": cv_0_01,
            "T": T,
            "N": N,
            "note": "Approximate critical values from Frees (1995)",
        }

    def _prepare_residual_data(self) -> pd.DataFrame:
        """
        Prepare residual data with entity and time identifiers.

        Returns
        -------
        pd.DataFrame
            DataFrame with columns: entity, time, resid
        """
        if hasattr(self.results, "entity_index") and hasattr(self.results, "time_index"):
            resid_flat = self.resid.ravel() if hasattr(self.resid, "ravel") else self.resid

            resid_df = pd.DataFrame(
                {
                    "entity": self.results.entity_index,
                    "time": self.results.time_index,
                    "resid": resid_flat,
                }
            )

            return resid_df
        else:
            raise AttributeError(
                "Results object must have 'entity_index' and 'time_index' attributes"
            )
