"""
Utility functions for unit root testing with non-stationary volatility.

This module implements auxiliary functions for data preprocessing, detrending,
and lag selection as described in Cavaliere & Taylor (2007, 2008).
"""

import numpy as np
from typing import Tuple, Optional, Union
from scipy import stats


def detrend(
    X: np.ndarray,
    trend: str = "c",
    return_params: bool = False
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """
    OLS detrending of time series.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data (T+1,) array starting from t=0
    trend : str, default="c"
        Type of deterministic component:
        - "c": constant only
        - "ct": constant and linear trend
        - "ctt": constant, linear and quadratic trend
    return_params : bool, default=False
        If True, return both residuals and estimated parameters
        
    Returns
    -------
    X_detrended : np.ndarray
        Detrended series
    gamma : np.ndarray, optional
        Estimated deterministic parameters (if return_params=True)
        
    Notes
    -----
    Implements standard OLS detrending as in Cavaliere & Taylor (2007, Section 3).
    """
    T = len(X) - 1  # Effective sample size (excluding t=0)
    
    # Construct deterministic regressor matrix
    if trend == "c":
        Z = np.ones((T + 1, 1))
    elif trend == "ct":
        t = np.arange(T + 1)
        Z = np.column_stack([np.ones(T + 1), t])
    elif trend == "ctt":
        t = np.arange(T + 1)
        Z = np.column_stack([np.ones(T + 1), t, t**2])
    else:
        raise ValueError(f"Unknown trend specification: {trend}")
    
    # OLS estimation: gamma = (Z'Z)^{-1} Z'X
    gamma = np.linalg.solve(Z.T @ Z, Z.T @ X)
    
    # Compute residuals
    X_detrended = X - Z @ gamma
    
    if return_params:
        return X_detrended, gamma
    return X_detrended


def local_gls_detrend(
    X: np.ndarray,
    trend: str = "c",
    c: float = None,
    return_params: bool = False
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """
    Local GLS detrending as in Elliott, Rothenberg & Stock (1996).
    
    Parameters
    ----------
    X : np.ndarray
        Time series data (T+1,) array starting from t=0
    trend : str, default="c"
        Type of deterministic component ("c" or "ct")
    c : float, optional
        Local-to-unity parameter. If None, uses:
        - c=7 for trend="c" (at 5% level)
        - c=13.5 for trend="ct" (at 5% level)
    return_params : bool, default=False
        If True, return both residuals and estimated parameters
        
    Returns
    -------
    X_detrended : np.ndarray
        Locally GLS detrended series
    gamma : np.ndarray, optional
        Estimated deterministic parameters (if return_params=True)
        
    Notes
    -----
    Implements local GLS detrending as described in Elliott et al. (1996)
    and used in Cavaliere & Taylor (2008, Section 4.1).
    
    References
    ----------
    Elliott, G., Rothenberg, T.J., Stock, J.H. (1996). "Efficient Tests for an 
    Autoregressive Unit Root". Econometrica 64, 813-836.
    """
    T = len(X) - 1
    
    # Set default c values based on trend
    if c is None:
        c = 7.0 if trend == "c" else 13.5
    
    alpha_bar = 1 - c / T
    
    # Construct quasi-differenced data
    X_qd = np.zeros(T + 1)
    X_qd[0] = X[0]
    X_qd[1:] = X[1:] - alpha_bar * X[:-1]
    
    # Construct quasi-differenced regressors
    if trend == "c":
        Z_qd = np.ones((T + 1, 1))
        Z_qd[1:, 0] = 1 - alpha_bar
    elif trend == "ct":
        t = np.arange(T + 1)
        Z_qd = np.column_stack([np.ones(T + 1), t])
        Z_qd[1:, 0] = 1 - alpha_bar
        Z_qd[1:, 1] = t[1:] - alpha_bar * t[:-1]
    else:
        raise ValueError(f"Local GLS detrending only supports 'c' or 'ct', got {trend}")
    
    # GLS estimation
    gamma = np.linalg.solve(Z_qd.T @ Z_qd, Z_qd.T @ X_qd)
    
    # Compute detrended series
    if trend == "c":
        Z = np.ones((T + 1, 1))
    else:  # trend == "ct"
        t = np.arange(T + 1)
        Z = np.column_stack([np.ones(T + 1), t])
    
    X_detrended = X - Z @ gamma
    
    if return_params:
        return X_detrended, gamma
    return X_detrended


def ar_spectral_density(
    X_detrended: np.ndarray,
    k: int,
    demean: bool = False
) -> float:
    """
    Autoregressive spectral density estimator at frequency zero.
    
    Parameters
    ----------
    X_detrended : np.ndarray
        Detrended time series (T+1,) array
    k : int
        Lag truncation parameter
    demean : bool, default=False
        Whether to demean the residuals
        
    Returns
    -------
    s2_AR : float
        AR spectral density estimate
        
    Notes
    -----
    Implements the AR spectral density estimator s²_AR(k) as in equation (5)
    of Cavaliere & Taylor (2008):
    
    s²_AR(k) = σ̂² / (1 - β̂(1))²
    
    where β̂(1) = Σᵢ₌₁ᵏ β̂ᵢ from the ADF regression:
    
    ΔX̂ₜ = ρX̂ₜ₋₁ + Σᵢ₌₁ᵏ βᵢΔX̂ₜ₋ᵢ + εₜ
    """
    T = len(X_detrended) - 1
    
    # Compute first differences
    dX = np.diff(X_detrended)
    
    if k == 0:
        # Simple variance estimator when k=0
        sigma2 = np.var(dX, ddof=1)
        return sigma2
    
    # Construct ADF regression
    y = dX[k:]  # ΔX_t for t = k+1, ..., T
    X_lag = X_detrended[k:-1]  # X_{t-1} for t = k+1, ..., T
    
    # Lagged differences: ΔX_{t-1}, ..., ΔX_{t-k}
    dX_lags = np.column_stack([dX[k-i-1:-i-1] for i in range(k)])
    
    # Combine regressors: [X_{t-1}, ΔX_{t-1}, ..., ΔX_{t-k}]
    regressors = np.column_stack([X_lag, dX_lags])
    
    # OLS estimation
    params = np.linalg.lstsq(regressors, y, rcond=None)[0]
    residuals = y - regressors @ params
    
    # Extract parameters
    rho = params[0]
    beta = params[1:]  # β₁, ..., βₖ
    
    # Compute σ̂²
    sigma2 = np.var(residuals, ddof=k+1) if demean else np.mean(residuals**2)
    
    # Compute β̂(1) = Σᵢ₌₁ᵏ β̂ᵢ
    beta_sum = np.sum(beta)
    
    # AR spectral density: s²_AR(k) = σ̂² / (1 - β̂(1))²
    s2_AR = sigma2 / (1 - beta_sum)**2
    
    return s2_AR


def select_lag_length(
    X_detrended: np.ndarray,
    max_lag: int = None,
    criterion: str = "MAIC"
) -> int:
    """
    Data-dependent lag length selection.
    
    Parameters
    ----------
    X_detrended : np.ndarray
        Detrended time series
    max_lag : int, optional
        Maximum lag to consider. If None, uses ⌊12(T/100)^{1/4}⌋
    criterion : str, default="MAIC"
        Information criterion:
        - "MAIC": Modified AIC (Ng & Perron, 2001)
        - "AIC": Akaike Information Criterion
        - "BIC": Bayesian Information Criterion
        
    Returns
    -------
    k_opt : int
        Optimal lag length
        
    Notes
    -----
    Implements MAIC criterion of Ng & Perron (2001) as used in 
    Cavaliere & Taylor (2008, Section 5).
    
    References
    ----------
    Ng, S., Perron, P. (2001). "Lag Length Selection and the Construction of 
    Unit Root Tests with Good Size and Power". Econometrica 69, 1519-1554.
    """
    T = len(X_detrended) - 1
    
    if max_lag is None:
        max_lag = int(np.floor(12 * (T / 100) ** 0.25))
    
    # Ensure max_lag is reasonable
    max_lag = min(max_lag, T // 3)
    
    # Compute first differences
    dX = np.diff(X_detrended)
    
    # Initialize criterion values
    ic_values = np.zeros(max_lag + 1)
    
    for k in range(max_lag + 1):
        if k == 0:
            # Simple AR(0) model
            sigma2 = np.var(dX, ddof=1)
            n_params = 1
        else:
            # ADF regression with k lags
            y = dX[k:]
            X_lag = X_detrended[k:-1]
            dX_lags = np.column_stack([dX[k-i-1:-i-1] for i in range(k)])
            regressors = np.column_stack([X_lag, dX_lags])
            
            params = np.linalg.lstsq(regressors, y, rcond=None)[0]
            residuals = y - regressors @ params
            sigma2 = np.var(residuals, ddof=k+1)
            n_params = k + 1
        
        # Compute information criterion
        n = T - k  # Effective sample size
        
        if criterion == "AIC":
            ic_values[k] = np.log(sigma2) + 2 * n_params / n
        elif criterion == "BIC":
            ic_values[k] = np.log(sigma2) + n_params * np.log(n) / n
        elif criterion == "MAIC":
            # Modified AIC: MAIC(k) = ln(σ̂²(k)) + 2(τ̂²(k) + k·σ̂²(k))/(T - k)
            # where τ̂²(k) is based on AR spectral density
            if k > 0:
                s2_AR = ar_spectral_density(X_detrended, k)
                tau2 = s2_AR - sigma2
                ic_values[k] = np.log(sigma2) + 2 * (tau2 + k * sigma2) / n
            else:
                ic_values[k] = np.log(sigma2)
        else:
            raise ValueError(f"Unknown criterion: {criterion}")
    
    k_opt = np.argmin(ic_values)
    return k_opt


def estimate_residuals(
    X_detrended: np.ndarray,
    method: str = "diff"
) -> np.ndarray:
    """
    Estimate residuals for variance profile estimation.
    
    Parameters
    ----------
    X_detrended : np.ndarray
        Detrended time series (T+1,) array
    method : str, default="diff"
        Residual estimation method:
        - "diff": First differences (imposing unit root)
        - "ar1": AR(1) residuals without imposing unit root
        - "adf": ADF regression residuals
        
    Returns
    -------
    residuals : np.ndarray
        Estimated residuals (T,) array
        
    Notes
    -----
    As discussed in Cavaliere & Taylor (2007, Remark 4.1), multiple
    methods can be used to estimate residuals for the variance profile.
    The "diff" method performs best in finite samples.
    """
    T = len(X_detrended) - 1
    
    if method == "diff":
        # Simple first differences: u_t = ΔX_t
        residuals = np.diff(X_detrended)
        
    elif method == "ar1":
        # AR(1) residuals: regress X_t on X_{t-1}
        y = X_detrended[1:]
        X_lag = X_detrended[:-1]
        
        alpha = np.sum(X_lag * y) / np.sum(X_lag**2)
        residuals = y - alpha * X_lag
        
    elif method == "adf":
        # This would require specifying k; default to diff
        residuals = np.diff(X_detrended)
        
    else:
        raise ValueError(f"Unknown residual method: {method}")
    
    return residuals


def check_assumptions(
    X: np.ndarray,
    max_lag: int = None,
    alpha: float = 0.05
) -> dict:
    """
    Check key assumptions for unit root tests.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    max_lag : int, optional
        Maximum lag for diagnostic tests
    alpha : float, default=0.05
        Significance level for tests
        
    Returns
    -------
    results : dict
        Dictionary containing diagnostic test results
    """
    T = len(X) - 1
    
    if max_lag is None:
        max_lag = int(np.floor(12 * (T / 100) ** 0.25))
    
    results = {
        "sample_size": T,
        "mean": np.mean(X),
        "variance": np.var(X),
        "autocorrelations": {},
        "warnings": []
    }
    
    # Compute autocorrelations
    dX = np.diff(X)
    for lag in range(1, min(max_lag + 1, T // 4)):
        acf = np.corrcoef(dX[lag:], dX[:-lag])[0, 1]
        results["autocorrelations"][f"lag_{lag}"] = acf
    
    # Check for explosiveness
    if np.any(np.abs(np.diff(X)) > 10 * np.std(np.diff(X))):
        results["warnings"].append("Potential explosive behavior detected")
    
    # Check sample size
    if T < 50:
        results["warnings"].append(f"Small sample size (T={T}). Results may be unreliable.")
    
    return results
