"""
Unit Root Tests Module

This module provides unit root tests commonly used for determining the 
order of integration of time series data before applying asymmetric 
causality tests.

Includes:
    - Augmented Dickey-Fuller (ADF) test
    - Phillips-Perron (PP) test
    - KPSS test
"""

import numpy as np
from scipy import stats
from scipy.stats import t as t_dist
from typing import Tuple, Optional, Dict
from dataclasses import dataclass


@dataclass
class UnitRootResult:
    """
    Data class for unit root test results.
    
    Attributes
    ----------
    test_name : str
        Name of the unit root test.
    statistic : float
        Test statistic value.
    p_value : float
        P-value of the test.
    critical_values : Dict[str, float]
        Critical values at standard significance levels.
    lags : int
        Number of lags used in the test.
    conclusion : str
        Conclusion about stationarity.
    """
    test_name: str
    statistic: float
    p_value: float
    critical_values: Dict[str, float]
    lags: int
    conclusion: str
    
    def __repr__(self) -> str:
        return (f"UnitRootResult(test='{self.test_name}', "
                f"statistic={self.statistic:.4f}, p_value={self.p_value:.4f})")
    
    def summary(self) -> str:
        """Generate formatted summary."""
        output = f"""
{self.test_name}
{'=' * 60}
Test Statistic:      {self.statistic:.4f}
P-value:             {self.p_value:.4f}
Lags Used:           {self.lags}

Critical Values:
  1%:   {self.critical_values.get('1%', 'N/A')}
  5%:   {self.critical_values.get('5%', 'N/A')}
  10%:  {self.critical_values.get('10%', 'N/A')}

Conclusion: {self.conclusion}
{'=' * 60}
"""
        return output


def _adf_critical_values(n: int, trend: str = 'c') -> Dict[str, float]:
    """
    Get critical values for ADF test.
    
    Based on MacKinnon (1996) response surface regressions.
    """
    # Approximate critical values
    if trend == 'n':  # No constant, no trend
        return {'1%': -2.58, '5%': -1.95, '10%': -1.62}
    elif trend == 'c':  # Constant only
        return {'1%': -3.43, '5%': -2.86, '10%': -2.57}
    elif trend == 'ct':  # Constant and trend
        return {'1%': -3.96, '5%': -3.41, '10%': -3.13}
    else:
        return {'1%': -3.43, '5%': -2.86, '10%': -2.57}


def _select_lag_aic(y: np.ndarray, max_lag: int, trend: str = 'c') -> int:
    """Select lag order using AIC criterion."""
    n = len(y)
    dy = np.diff(y)
    
    best_aic = np.inf
    best_lag = 0
    
    for lag in range(0, min(max_lag + 1, len(dy) // 3)):
        # Construct regression matrices
        effective_n = len(dy) - lag
        if effective_n < 10:
            continue
            
        Y = dy[lag:]
        X = y[lag:-1].reshape(-1, 1)
        
        # Add trend components
        if trend in ['c', 'ct']:
            X = np.column_stack([np.ones(effective_n), X])
        if trend == 'ct':
            X = np.column_stack([X, np.arange(1, effective_n + 1)])
        
        # Add lagged differences
        for i in range(1, lag + 1):
            X = np.column_stack([X, dy[lag - i:-(i) if i < lag else None][:effective_n]])
        
        try:
            beta = np.linalg.lstsq(X, Y, rcond=None)[0]
            resid = Y - X @ beta
            sse = np.sum(resid**2)
            k = X.shape[1]
            aic = effective_n * np.log(sse / effective_n) + 2 * k
            
            if aic < best_aic:
                best_aic = aic
                best_lag = lag
        except:
            continue
    
    return best_lag


def adf_test(y: np.ndarray, 
             max_lag: Optional[int] = None,
             trend: str = 'c',
             autolag: bool = True) -> UnitRootResult:
    """
    Augmented Dickey-Fuller test for unit root.
    
    Parameters
    ----------
    y : np.ndarray
        Time series data.
    max_lag : int, optional
        Maximum lag order. If None, uses int(12*(T/100)^(1/4)).
    trend : str, optional
        Trend specification:
        - 'n': No constant, no trend
        - 'c': Constant only (default)
        - 'ct': Constant and trend
    autolag : bool, optional
        If True, automatically select lag order using AIC.
    
    Returns
    -------
    UnitRootResult
        Test results with statistic, p-value, and conclusion.
    
    Notes
    -----
    H0: The series has a unit root (non-stationary).
    H1: The series is stationary.
    """
    y = np.asarray(y, dtype=np.float64).flatten()
    n = len(y)
    
    # Default max lag (Schwert, 1989)
    if max_lag is None:
        max_lag = int(12 * (n / 100) ** 0.25)
    
    # Select lag order
    if autolag:
        lag = _select_lag_aic(y, max_lag, trend)
    else:
        lag = max_lag
    
    # Calculate differences
    dy = np.diff(y)
    
    # Effective sample size
    effective_n = len(dy) - lag
    
    # Dependent variable
    Y = dy[lag:]
    
    # Level variable
    X = y[lag:-1].reshape(-1, 1)
    
    # Add deterministic components
    if trend in ['c', 'ct']:
        X = np.column_stack([np.ones(effective_n), X])
    if trend == 'ct':
        X = np.column_stack([X, np.arange(1, effective_n + 1)])
    
    # Add lagged differences
    for i in range(1, lag + 1):
        start_idx = lag - i
        end_idx = -i if i < lag else len(dy) - lag
        lagged_diff = dy[start_idx:start_idx + effective_n]
        X = np.column_stack([X, lagged_diff])
    
    # OLS estimation
    try:
        XtX_inv = np.linalg.inv(X.T @ X)
        beta = XtX_inv @ X.T @ Y
        resid = Y - X @ beta
        sse = np.sum(resid**2)
        s2 = sse / (effective_n - X.shape[1])
        
        # Standard error of coefficient on lagged level
        if trend == 'n':
            rho_idx = 0
        else:
            rho_idx = 1  # After constant
        
        se_rho = np.sqrt(s2 * XtX_inv[rho_idx, rho_idx])
        rho = beta[rho_idx]
        
        # ADF statistic
        adf_stat = rho / se_rho
        
    except np.linalg.LinAlgError:
        return UnitRootResult(
            test_name="Augmented Dickey-Fuller Test",
            statistic=np.nan,
            p_value=np.nan,
            critical_values=_adf_critical_values(n, trend),
            lags=lag,
            conclusion="Could not compute test (singular matrix)"
        )
    
    # Critical values
    cvs = _adf_critical_values(n, trend)
    
    # Approximate p-value using MacKinnon (1994) approach
    # This is a simplified approximation
    if adf_stat < cvs['1%']:
        p_value = 0.005
    elif adf_stat < cvs['5%']:
        p_value = 0.025
    elif adf_stat < cvs['10%']:
        p_value = 0.075
    else:
        p_value = 0.5 + 0.5 * (1 - stats.norm.cdf(-adf_stat))
    
    p_value = min(max(p_value, 0.0001), 0.9999)
    
    # Conclusion
    if adf_stat < cvs['5%']:
        conclusion = "Reject H0: Series appears stationary (I(0))"
    else:
        conclusion = "Fail to reject H0: Series appears non-stationary (has unit root)"
    
    return UnitRootResult(
        test_name="Augmented Dickey-Fuller Test",
        statistic=float(adf_stat),
        p_value=float(p_value),
        critical_values=cvs,
        lags=lag,
        conclusion=conclusion
    )


def kpss_test(y: np.ndarray,
              trend: str = 'c',
              lags: Optional[int] = None) -> UnitRootResult:
    """
    KPSS test for stationarity.
    
    Parameters
    ----------
    y : np.ndarray
        Time series data.
    trend : str, optional
        Trend specification:
        - 'c': Level stationarity (default)
        - 'ct': Trend stationarity
    lags : int, optional
        Number of lags for HAC variance. If None, uses automatic selection.
    
    Returns
    -------
    UnitRootResult
        Test results.
    
    Notes
    -----
    H0: The series is stationary.
    H1: The series has a unit root.
    
    Note: This is opposite to ADF where H0 is unit root.
    """
    y = np.asarray(y, dtype=np.float64).flatten()
    n = len(y)
    
    # Default lags (Newey-West)
    if lags is None:
        lags = int(4 * (n / 100) ** 0.25)
    
    # Detrending
    if trend == 'c':
        X = np.ones((n, 1))
    else:  # ct
        X = np.column_stack([np.ones(n), np.arange(1, n + 1)])
    
    beta = np.linalg.lstsq(X, y, rcond=None)[0]
    resid = y - X @ beta
    
    # Cumulative sum of residuals
    S = np.cumsum(resid)
    
    # HAC variance estimation (Newey-West)
    gamma0 = np.sum(resid**2) / n
    
    s2 = gamma0
    for j in range(1, lags + 1):
        weight = 1 - j / (lags + 1)  # Bartlett kernel
        gamma_j = np.sum(resid[j:] * resid[:-j]) / n
        s2 += 2 * weight * gamma_j
    
    # KPSS statistic
    kpss_stat = np.sum(S**2) / (n**2 * s2)
    
    # Critical values
    if trend == 'c':
        cvs = {'1%': 0.739, '5%': 0.463, '10%': 0.347}
    else:  # ct
        cvs = {'1%': 0.216, '5%': 0.146, '10%': 0.119}
    
    # Approximate p-value
    if kpss_stat > cvs['1%']:
        p_value = 0.005
    elif kpss_stat > cvs['5%']:
        p_value = 0.025
    elif kpss_stat > cvs['10%']:
        p_value = 0.075
    else:
        p_value = 0.15
    
    # Conclusion
    if kpss_stat > cvs['5%']:
        conclusion = "Reject H0: Series appears non-stationary"
    else:
        conclusion = "Fail to reject H0: Series appears stationary"
    
    return UnitRootResult(
        test_name="KPSS Test",
        statistic=float(kpss_stat),
        p_value=float(p_value),
        critical_values=cvs,
        lags=lags,
        conclusion=conclusion
    )


def pp_test(y: np.ndarray,
            trend: str = 'c',
            lags: Optional[int] = None) -> UnitRootResult:
    """
    Phillips-Perron test for unit root.
    
    Parameters
    ----------
    y : np.ndarray
        Time series data.
    trend : str, optional
        Trend specification ('n', 'c', 'ct').
    lags : int, optional
        Number of lags for HAC variance.
    
    Returns
    -------
    UnitRootResult
        Test results.
    
    Notes
    -----
    H0: The series has a unit root.
    H1: The series is stationary.
    """
    y = np.asarray(y, dtype=np.float64).flatten()
    n = len(y)
    
    # Default lags
    if lags is None:
        lags = int(4 * (n / 100) ** 0.25)
    
    # Construct regression
    Y = y[1:]
    X = y[:-1].reshape(-1, 1)
    
    if trend in ['c', 'ct']:
        X = np.column_stack([np.ones(n - 1), X])
    if trend == 'ct':
        X = np.column_stack([X, np.arange(1, n)])
    
    # OLS estimation
    beta = np.linalg.lstsq(X, Y, rcond=None)[0]
    resid = Y - X @ beta
    
    # HAC variance
    gamma0 = np.sum(resid**2) / n
    
    lambda_sq = gamma0
    for j in range(1, lags + 1):
        weight = 1 - j / (lags + 1)
        gamma_j = np.sum(resid[j:] * resid[:-j]) / n
        lambda_sq += 2 * weight * gamma_j
    
    # Standard error
    XtX_inv = np.linalg.inv(X.T @ X)
    s2 = np.sum(resid**2) / (n - 1 - X.shape[1])
    
    rho_idx = 1 if trend in ['c', 'ct'] else 0
    se_rho = np.sqrt(s2 * XtX_inv[rho_idx, rho_idx])
    rho = beta[rho_idx]
    
    # PP adjustment
    pp_stat = (np.sqrt(gamma0 / lambda_sq) * (rho - 1) / se_rho - 
               0.5 * (lambda_sq - gamma0) * np.sqrt(n) / (np.sqrt(lambda_sq) * np.sqrt(s2) * 
               np.sqrt(np.sum((X[:, rho_idx] - X[:, rho_idx].mean())**2))))
    
    # Critical values (same as ADF)
    cvs = _adf_critical_values(n, trend)
    
    # Approximate p-value
    if pp_stat < cvs['1%']:
        p_value = 0.005
    elif pp_stat < cvs['5%']:
        p_value = 0.025
    elif pp_stat < cvs['10%']:
        p_value = 0.075
    else:
        p_value = 0.5
    
    # Conclusion
    if pp_stat < cvs['5%']:
        conclusion = "Reject H0: Series appears stationary"
    else:
        conclusion = "Fail to reject H0: Series appears non-stationary (has unit root)"
    
    return UnitRootResult(
        test_name="Phillips-Perron Test",
        statistic=float(pp_stat),
        p_value=float(p_value),
        critical_values=cvs,
        lags=lags,
        conclusion=conclusion
    )


def determine_integration_order(y: np.ndarray,
                                max_order: int = 2,
                                significance: float = 0.05) -> Tuple[int, Dict]:
    """
    Determine the order of integration of a time series.
    
    Parameters
    ----------
    y : np.ndarray
        Time series data.
    max_order : int, optional
        Maximum integration order to test. Default is 2.
    significance : float, optional
        Significance level for ADF test. Default is 0.05.
    
    Returns
    -------
    Tuple[int, Dict]
        Tuple containing:
        - integration_order: Estimated order of integration (0, 1, or 2)
        - test_results: Dictionary of ADF test results at each level
    
    Notes
    -----
    This function performs sequential ADF tests to determine if the series
    is I(0), I(1), or I(2).
    """
    y = np.asarray(y, dtype=np.float64).flatten()
    
    test_results = {}
    current_series = y.copy()
    
    for order in range(max_order + 1):
        if order > 0:
            current_series = np.diff(current_series)
        
        if len(current_series) < 20:
            return order - 1, test_results
        
        result = adf_test(current_series, trend='c')
        test_results[f'd{order}'] = result
        
        # Check if series is stationary at this order
        cv_5pct = result.critical_values.get('5%', -2.86)
        
        if result.statistic < cv_5pct:
            return order, test_results
    
    # If still non-stationary after max differencing
    return max_order, test_results


def print_unit_root_summary(results: Dict[str, UnitRootResult]) -> str:
    """
    Generate formatted summary of unit root tests.
    
    Parameters
    ----------
    results : Dict[str, UnitRootResult]
        Dictionary of unit root test results.
    
    Returns
    -------
    str
        Formatted summary string.
    """
    separator = "=" * 70
    output = f"""
{separator}
                    UNIT ROOT TEST RESULTS
{separator}

"""
    
    for name, result in results.items():
        output += result.summary() + "\n"
    
    return output
