"""
Simulation-Based Unit Root Tests for Non-Stationary Volatility.

Implements the simulation-based methodology from Cavaliere & Taylor (2007),
"Testing for unit roots in time series models with non-stationary volatility",
Journal of Econometrics, 140, 919-947.
"""

import numpy as np
from typing import Dict, Optional, Callable
from scipy.interpolate import interp1d
from .utils import detrend, estimate_residuals
from .m_tests import compute_m_tests


def estimate_variance_profile(
    X: np.ndarray,
    trend: str = "c",
    method: str = "diff",
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Estimate the variance profile η(s).
    
    Parameters
    ----------
    X : np.ndarray
        Time series data (T+1,) array
    trend : str, default="c"
        Deterministic trend for detrending
    method : str, default="diff"
        Residual estimation method: "diff", "ar1", "adf"
        
    Returns
    -------
    eta_hat : callable
        Estimated variance profile function η̂(s) on [0,1]
        
    Notes
    -----
    Implements the variance profile estimator from Cavaliere & Taylor (2007, Section 4.1).
    
    The variance profile is defined (equation 4) as:
    
    η(s) = (∫₀¹ ω(r)² dr)⁻¹ ∫₀ˢ ω(r)² dr
    
    The estimator (equation 8) is:
    
    η̂(s) = Σᵗ₌₁⌊ˢᵀ⌋ û²ₜ / Σᵗ₌₁ᵀ û²ₜ
    
    with linear interpolation between sample points.
    
    Theorem 2 establishes that η̂(s) - η(s) →ₚ 0 uniformly for all s ∈ [0,1].
    
    References
    ----------
    Cavaliere, G., Taylor, A.M.R. (2007). "Testing for unit roots in time series 
    models with non-stationary volatility". Journal of Econometrics 140, 919-947.
    """
    T = len(X) - 1
    
    # Step 1: Detrend the data
    X_hat = detrend(X, trend=trend)
    
    # Step 2: Estimate residuals
    u_hat = estimate_residuals(X_hat, method=method)
    
    # Ensure we have T residuals
    if len(u_hat) < T:
        u_hat = np.concatenate([u_hat, [u_hat[-1]] * (T - len(u_hat))])
    u_hat = u_hat[:T]
    
    # Step 3: Compute squared residuals
    u_hat_sq = u_hat ** 2
    
    # Step 4: Compute cumulative sum of squared residuals
    cum_sum = np.cumsum(u_hat_sq)
    total_sum = cum_sum[-1]
    
    # Step 5: Construct variance profile estimator
    # η̂(s) for s = t/T, t = 0, 1, ..., T
    s_grid = np.linspace(0, 1, T + 1)
    eta_values = np.zeros(T + 1)
    eta_values[0] = 0.0
    eta_values[1:] = cum_sum / total_sum
    eta_values[-1] = 1.0  # Ensure η̂(1) = 1
    
    # Step 6: Create interpolated function with linear interpolation
    eta_hat = interp1d(
        s_grid, 
        eta_values, 
        kind='linear', 
        bounds_error=False, 
        fill_value=(0.0, 1.0)
    )
    
    return eta_hat


def simulate_variance_transformed_brownian(
    eta: Callable[[np.ndarray], np.ndarray],
    T: int = 1000,
    seed: Optional[int] = None,
) -> np.ndarray:
    """
    Simulate variance-transformed Brownian motion B_η(s).
    
    Parameters
    ----------
    eta : callable
        Variance profile function
    T : int, default=1000
        Number of time steps
    seed : int, optional
        Random seed
        
    Returns
    -------
    B_eta : np.ndarray
        Simulated variance-transformed Brownian motion on grid [0, 1/T, 2/T, ..., 1]
        
    Notes
    -----
    Implements the approximation from Cavaliere & Taylor (2007, Theorem 3):
    
    B_η(s) ≈ T^{-1/2} Σᵗ₌₁⌊η̂(⌊ᵀˢ⌋/T)T⌋ εₜ
    
    where εₜ ~ iid N(0,1).
    """
    if seed is not None:
        np.random.seed(seed)
    
    # Generate iid standard normal innovations
    epsilon = np.random.normal(0, 1, T)
    
    # Create grid on [0, 1]
    s_grid = np.linspace(0, 1, T + 1)
    
    # Initialize Brownian motion
    B_eta = np.zeros(T + 1)
    
    # For each s in grid, compute B_η(s)
    for i, s in enumerate(s_grid[1:], 1):
        # Compute η(s)
        eta_s = eta(s)
        
        # Number of terms to sum: ⌊η(s) × T⌋
        n_terms = int(np.floor(eta_s * T))
        n_terms = min(n_terms, T)  # Ensure we don't exceed T
        
        # B_η(s) = T^{-1/2} Σᵗ₌₁ⁿ εₜ
        if n_terms > 0:
            B_eta[i] = np.sum(epsilon[:n_terms]) / np.sqrt(T)
    
    return B_eta


def simulate_critical_values(
    eta: Callable[[np.ndarray], np.ndarray],
    statistic: str = "MZa",
    trend: str = "c",
    N: int = 10000,
    T: int = 1000,
    alpha_levels: Optional[np.ndarray] = None,
    seed: Optional[int] = None,
) -> Dict[float, float]:
    """
    Simulate critical values for M tests under non-stationary volatility.
    
    Parameters
    ----------
    eta : callable
        Variance profile function
    statistic : str, default="MZa"
        Test statistic: "MZa", "MZt", or "MSB"
    trend : str, default="c"
        Deterministic trend
    N : int, default=10000
        Number of simulation replications
    T : int, default=1000
        Number of time steps for Brownian motion simulation
    alpha_levels : np.ndarray, optional
        Significance levels. Default: [0.01, 0.05, 0.10]
    seed : int, optional
        Random seed
        
    Returns
    -------
    critical_values : dict
        Dictionary mapping significance levels to critical values
        
    Notes
    -----
    Implements the simulation procedure from Cavaliere & Taylor (2007, Section 4.2).
    
    Algorithm:
    1. Generate B_η,T(s) = T^{-1/2} Σᵗ₌₁⌊η̂(⌊ᵀˢ⌋/T)T⌋ εₜ, εₜ ~ iid N(0,1)
    2. Construct de-trended process B̃_η|Z by OLS projection
    3. Compute test statistic functional
    4. Repeat N times to obtain distribution
    5. Extract quantiles
    """
    if seed is not None:
        np.random.seed(seed)
    
    if alpha_levels is None:
        alpha_levels = np.array([0.01, 0.05, 0.10])
    
    # Storage for simulated statistics
    simulated_stats = np.zeros(N)
    
    for i in range(N):
        # Step 1: Simulate variance-transformed Brownian motion
        B_eta_T = simulate_variance_transformed_brownian(eta, T=T)
        
        # Step 2: Create deterministic regressors
        s_grid = np.linspace(0, 1, T + 1)
        
        if trend == "c":
            Z = np.ones((T + 1, 1))
        elif trend == "ct":
            Z = np.column_stack([np.ones(T + 1), s_grid])
        else:
            raise ValueError(f"Unknown trend: {trend}")
        
        # Step 3: OLS projection to get B̃_η|Z
        # B̃_η|Z(s) = B_η(s) - Z(s)'(∫Z(r)Z(r)'dr)^{-1}∫Z(r)B_η(r)dr
        ZZ_inv = np.linalg.inv(Z.T @ Z / (T + 1))
        ZB = Z.T @ B_eta_T / (T + 1)
        projection = Z @ ZZ_inv @ ZB
        
        B_tilde = B_eta_T - projection
        
        # Step 4: Compute test statistic functional
        if statistic == "MZa":
            # MZa limit: 0.5(B̃²(1) - B̃²(0) - 1) / ∫B̃²(s)ds
            numerator = 0.5 * (B_tilde[-1]**2 - B_tilde[0]**2 - 1)
            denominator = np.trapz(B_tilde**2, s_grid)
            simulated_stats[i] = numerator / denominator if denominator > 0 else -np.inf
            
        elif statistic == "MZt":
            # MZt limit: 0.5(B̃²(1) - B̃²(0) - 1) / (∫B̃²(s)ds)^{1/2}
            numerator = 0.5 * (B_tilde[-1]**2 - B_tilde[0]**2 - 1)
            denominator = np.sqrt(np.trapz(B_tilde**2, s_grid))
            simulated_stats[i] = numerator / denominator if denominator > 0 else -np.inf
            
        elif statistic == "MSB":
            # MSB limit: (∫B̃²(s)ds)^{1/2}
            simulated_stats[i] = np.sqrt(np.trapz(B_tilde**2, s_grid))
            
        else:
            raise ValueError(f"Unknown statistic: {statistic}")
    
    # Step 5: Compute critical values (left-tail quantiles)
    critical_values = {}
    for alpha in alpha_levels:
        critical_values[alpha] = np.percentile(simulated_stats, alpha * 100)
    
    return critical_values


def simulation_based_test(
    X: np.ndarray,
    statistic: str = "MZa",
    trend: str = "c",
    k: Optional[int] = None,
    N: int = 10000,
    T_sim: int = 1000,
    alpha: float = 0.05,
    seed: Optional[int] = None,
) -> Dict[str, any]:
    """
    Simulation-based unit root test for non-stationary volatility.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data (T+1,) array
    statistic : str, default="MZa"
        Test statistic: "MZa", "MZt", or "MSB"
    trend : str, default="c"
        Deterministic trend
    k : int, optional
        Lag truncation parameter
    N : int, default=10000
        Number of simulation replications for critical values
    T_sim : int, default=1000
        Grid size for Brownian motion simulation
    alpha : float, default=0.05
        Significance level
    seed : int, optional
        Random seed
        
    Returns
    -------
    results : dict
        Dictionary containing:
        - "statistic": observed test statistic
        - "critical_value": simulated critical value at level alpha
        - "reject": rejection decision
        - "variance_profile": estimated variance profile function
        - "all_critical_values": critical values at [0.01, 0.05, 0.10]
        
    Notes
    -----
    Implements the complete simulation-based test procedure from 
    Cavaliere & Taylor (2007, Section 4.2).
    
    The test is asymptotically similar under the null hypothesis for any
    volatility process satisfying Assumption A3 (bounded with countable jumps).
    """
    if seed is not None:
        np.random.seed(seed)
    
    # Step 1: Compute observed test statistic
    obs_stats = compute_m_tests(X, trend=trend, k=k)
    obs_statistic = obs_stats[statistic]
    
    # Step 2: Estimate variance profile
    eta_hat = estimate_variance_profile(X, trend=trend)
    
    # Step 3: Simulate critical values using estimated variance profile
    alpha_levels = np.array([0.01, 0.05, 0.10])
    critical_values = simulate_critical_values(
        eta_hat,
        statistic=statistic,
        trend=trend,
        N=N,
        T=T_sim,
        alpha_levels=alpha_levels,
        seed=seed,
    )
    
    # Step 4: Make testing decision
    critical_value_alpha = critical_values[alpha]
    reject = obs_statistic < critical_value_alpha
    
    results = {
        "statistic": obs_statistic,
        "critical_value": critical_value_alpha,
        "critical_values": critical_values,
        "reject": reject,
        "alpha": alpha,
        "variance_profile": eta_hat,
        "test": statistic,
        "trend": trend,
        "k": obs_stats["k"],
        "N": N,
        "T_sim": T_sim,
    }
    
    return results
