"""
Wild Bootstrap Unit Root Tests for Non-Stationary Volatility.

Implements the wild bootstrap methodology from Cavaliere & Taylor (2008),
"Bootstrap Unit Root Tests for Time Series with Nonstationary Volatility",
Econometric Theory, 24, 43-71.
"""

import numpy as np
from typing import Dict, Optional, Tuple, Union
from .utils import detrend, local_gls_detrend, select_lag_length, estimate_residuals
from .m_tests import compute_m_tests


def wild_bootstrap_test(
    X: np.ndarray,
    statistic: str = "MZa",
    trend: str = "c",
    detrend_method: str = "gls",
    k: Optional[int] = None,
    c: Optional[float] = None,
    N: int = 999,
    seed: Optional[int] = None,
    residual_method: str = "diff",
    distribution: str = "gaussian",
) -> Dict[str, any]:
    """
    Wild bootstrap unit root test robust to non-stationary volatility.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data (T+1,) array including initial value
    statistic : str, default="MZa"
        Test statistic: "MZa", "MZt", or "MSB"
    trend : str, default="c"
        Deterministic component: "c" or "ct"
    detrend_method : str, default="gls"
        Detrending method: "ols" or "gls"
    k : int, optional
        Lag truncation. If None, selected automatically
    c : float, optional
        Local-to-unity parameter for GLS (default: 7 for "c", 13.5 for "ct")
    N : int, default=999
        Number of bootstrap replications
    seed : int, optional
        Random seed for reproducibility
    residual_method : str, default="diff"
        Method for residual estimation: "diff", "ar1", or "adf"
    distribution : str, default="gaussian"
        Bootstrap distribution: "gaussian" or "mammen"
        
    Returns
    -------
    results : dict
        Dictionary containing:
        - "statistic": observed test statistic
        - "p_value": bootstrap p-value
        - "critical_values": bootstrap critical values at [0.01, 0.05, 0.10]
        - "bootstrap_distribution": array of bootstrap statistics
        - "reject_01", "reject_05", "reject_10": rejection decisions
        
    Notes
    -----
    Implements the wild bootstrap procedure from Cavaliere & Taylor (2008, Section 4).
    
    The algorithm:
    1. Compute residuals ût from ADF regression on local GLS detrended data
    2. Generate bootstrap residuals: u*t = ût × wt, where wt ~ N(0,1) iid
    3. Construct bootstrap sample: X*t = X*_{t-1} + u*t, X*0 = 0
    4. Compute bootstrap statistic on X*
    5. Repeat N times to obtain bootstrap distribution
    6. p-value = proportion of bootstrap statistics ≤ observed statistic
    
    The wild bootstrap preserves heteroskedasticity because conditionally on ût,
    u*t is independent with variance û²t.
    
    References
    ----------
    Cavaliere, G., Taylor, A.M.R. (2008). "Bootstrap Unit Root Tests for Time Series 
    with Nonstationary Volatility". Econometric Theory 24, 43-71.
    """
    if seed is not None:
        np.random.seed(seed)
    
    T = len(X) - 1
    
    # Step 1: Compute observed test statistic
    obs_stats = compute_m_tests(
        X, trend=trend, detrend_method=detrend_method, k=k, c=c
    )
    obs_statistic = obs_stats[statistic]
    k_selected = obs_stats["k"]
    
    # Step 2: Detrend data for residual estimation
    if detrend_method == "gls":
        X_hat = local_gls_detrend(X, trend=trend, c=c)
    else:
        X_hat = detrend(X, trend=trend)
    
    # Step 3: Estimate residuals (Section 4.1, Remark 5)
    residuals = estimate_residuals(X_hat, method=residual_method)
    
    # Ensure residuals have length T
    if len(residuals) < T:
        # Pad with the last residual value if needed
        residuals = np.concatenate([residuals, [residuals[-1]] * (T - len(residuals))])
    residuals = residuals[:T]
    
    # Step 4: Generate bootstrap samples and compute statistics
    bootstrap_stats = np.zeros(N)
    
    for i in range(N):
        # Generate wild bootstrap weights
        if distribution == "gaussian":
            w = np.random.normal(0, 1, T)
        elif distribution == "mammen":
            # Mammen (1993) two-point distribution
            # P(w = -(√5-1)/2) = (√5+1)/(2√5), P(w = (√5+1)/2) = 1 - (√5+1)/(2√5)
            pi = (np.sqrt(5) + 1) / (2 * np.sqrt(5))
            u = np.random.uniform(0, 1, T)
            w = np.where(u < pi, -(np.sqrt(5) - 1) / 2, (np.sqrt(5) + 1) / 2)
        else:
            raise ValueError(f"Unknown distribution: {distribution}")
        
        # Generate bootstrap residuals: u*t = ût × wt (equation 14-15)
        u_star = residuals * w
        
        # Generate bootstrap sample (equation 16)
        X_star = np.zeros(T + 1)
        X_star[0] = 0  # Initial value
        X_star[1:] = np.cumsum(u_star)
        
        # Compute bootstrap statistic (equations 17-18)
        boot_stats = compute_m_tests(
            X_star,
            trend=trend,
            detrend_method=detrend_method,
            k=k_selected,  # Use same lag as observed data
            c=c,
        )
        bootstrap_stats[i] = boot_stats[statistic]
    
    # Step 5: Compute p-value and critical values
    # For MZa and MZt: left-tailed test (reject for large negative values)
    # For MSB: left-tailed test (reject for small values)
    if statistic in ["MZa", "MZt", "MSB"]:
        # All are left-tailed
        p_value = np.mean(bootstrap_stats <= obs_statistic)
        cv_01 = np.percentile(bootstrap_stats, 1)
        cv_05 = np.percentile(bootstrap_stats, 5)
        cv_10 = np.percentile(bootstrap_stats, 10)
    else:
        raise ValueError(f"Unknown statistic: {statistic}")
    
    # Rejection decisions
    reject_01 = obs_statistic < cv_01
    reject_05 = obs_statistic < cv_05
    reject_10 = obs_statistic < cv_10
    
    results = {
        "statistic": obs_statistic,
        "p_value": p_value,
        "critical_values": {
            "1%": cv_01,
            "5%": cv_05,
            "10%": cv_10,
        },
        "reject_01": reject_01,
        "reject_05": reject_05,
        "reject_10": reject_10,
        "bootstrap_distribution": bootstrap_stats,
        "N": N,
        "k": k_selected,
        "test": statistic,
        "trend": trend,
        "detrend_method": detrend_method,
    }
    
    return results


def bootstrap_mza(
    X: np.ndarray,
    trend: str = "c",
    N: int = 999,
    **kwargs
) -> Dict[str, any]:
    """
    Wild bootstrap MZa test.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend
    N : int, default=999
        Number of bootstrap replications
    **kwargs : additional arguments
        Passed to wild_bootstrap_test
        
    Returns
    -------
    results : dict
        Bootstrap test results
    """
    return wild_bootstrap_test(X, statistic="MZa", trend=trend, N=N, **kwargs)


def bootstrap_mzt(
    X: np.ndarray,
    trend: str = "c",
    N: int = 999,
    **kwargs
) -> Dict[str, any]:
    """
    Wild bootstrap MZt test.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend
    N : int, default=999
        Number of bootstrap replications
    **kwargs : additional arguments
        Passed to wild_bootstrap_test
        
    Returns
    -------
    results : dict
        Bootstrap test results
    """
    return wild_bootstrap_test(X, statistic="MZt", trend=trend, N=N, **kwargs)


def bootstrap_msb(
    X: np.ndarray,
    trend: str = "c",
    N: int = 999,
    **kwargs
) -> Dict[str, any]:
    """
    Wild bootstrap MSB test.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend
    N : int, default=999
        Number of bootstrap replications
    **kwargs : additional arguments
        Passed to wild_bootstrap_test
        
    Returns
    -------
    results : dict
        Bootstrap test results
    """
    return wild_bootstrap_test(X, statistic="MSB", trend=trend, N=N, **kwargs)
