"""
M-class unit root statistics.

Implements the MZa, MZt, and MSB statistics from Perron & Ng (1996), 
Stock (1999), and Ng & Perron (2001) as described in Cavaliere & Taylor (2007, 2008).
"""

import numpy as np
from typing import Dict, Tuple, Optional
from .utils import (
    detrend,
    local_gls_detrend,
    ar_spectral_density,
    select_lag_length,
)


def compute_m_tests(
    X: np.ndarray,
    trend: str = "c",
    detrend_method: str = "ols",
    k: Optional[int] = None,
    c: Optional[float] = None,
    lag_selection: str = "MAIC",
    max_lag: Optional[int] = None,
) -> Dict[str, float]:
    """
    Compute all M unit root statistics.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data (T+1,) array including t=0
    trend : str, default="c"
        Deterministic component: "c" (constant) or "ct" (constant + trend)
    detrend_method : str, default="ols"
        Detrending method: "ols" or "gls" (local GLS)
    k : int, optional
        Lag truncation parameter. If None, selected by lag_selection criterion
    c : float, optional
        Local-to-unity parameter for GLS detrending (only used if detrend_method="gls")
    lag_selection : str, default="MAIC"
        Criterion for automatic lag selection ("MAIC", "AIC", "BIC")
    max_lag : int, optional
        Maximum lag for lag selection
        
    Returns
    -------
    results : dict
        Dictionary containing:
        - "MZa": MZa statistic
        - "MZt": MZt statistic  
        - "MSB": MSB statistic
        - "k": lag truncation used
        - "s2_AR": long-run variance estimate
        
    Notes
    -----
    Implements equations (8)-(9) from Cavaliere & Taylor (2008):
    
    MZa = (T⁻¹X̂²_T - s²_AR(k)) / (2T⁻²Σₜ X̂²ₜ₋₁)
    
    MSB = (T⁻²Σₜ X̂²ₜ₋₁ / s²_AR(k))^{1/2}
    
    MZt = MZa × MSB
    
    where X̂ₜ are residuals from regression of Xₜ on zₜ.
    
    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.
    
    Perron, P., Ng, S. (1996). "Useful Modifications to Some Unit Root Tests with 
    Dependent Errors and Their Local Asymptotic Properties". Review of Economic 
    Studies 63, 435-463.
    """
    T = len(X) - 1
    
    # Step 1: Detrend the data
    if detrend_method == "ols":
        X_hat = detrend(X, trend=trend)
    elif detrend_method == "gls":
        X_hat = local_gls_detrend(X, trend=trend, c=c)
    else:
        raise ValueError(f"Unknown detrending method: {detrend_method}")
    
    # Step 2: Select lag length if not provided
    if k is None:
        k = select_lag_length(X_hat, max_lag=max_lag, criterion=lag_selection)
    
    # Ensure k satisfies Assumption K: 1/k + k³/T → 0
    k = min(k, int((T / 3) ** (1/3)))
    
    # Step 3: Compute long-run variance estimator s²_AR(k)
    s2_AR = ar_spectral_density(X_hat, k)
    
    # Step 4: Compute statistics components
    # Note: X_hat[0] corresponds to t=0, X_hat[T] to t=T
    X_hat_T = X_hat[T]  # X̂_T
    X_hat_0 = X_hat[0]  # X̂_0
    
    # T⁻¹ X̂²_T
    term1 = X_hat_T**2 / T
    
    # T⁻¹ X̂²_0  
    term2 = X_hat_0**2 / T
    
    # T⁻² Σₜ₌₁ᵀ X̂²ₜ₋₁
    sum_X2_lag = np.sum(X_hat[:-1]**2) / (T**2)
    
    # Step 5: Compute M statistics (equation 8)
    
    # MZa = (T⁻¹X̂²_T - T⁻¹X̂²_0 - s²_AR(k)) / (2T⁻²Σₜ X̂²ₜ₋₁)
    MZa_stat = (term1 - term2 - s2_AR) / (2 * sum_X2_lag)
    
    # MSB = (T⁻²Σₜ X̂²ₜ₋₁ / s²_AR(k))^{1/2}
    MSB_stat = np.sqrt(sum_X2_lag / s2_AR)
    
    # MZt = MZa × MSB
    MZt_stat = MZa_stat * MSB_stat
    
    results = {
        "MZa": MZa_stat,
        "MZt": MZt_stat,
        "MSB": MSB_stat,
        "k": k,
        "s2_AR": s2_AR,
        "T": T,
        "detrend_method": detrend_method,
        "trend": trend,
    }
    
    return results


def MZa(
    X: np.ndarray,
    trend: str = "c",
    detrend_method: str = "ols",
    k: Optional[int] = None,
    **kwargs
) -> float:
    """
    Compute MZa unit root statistic.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend specification
    detrend_method : str, default="ols"
        Detrending method
    k : int, optional
        Lag truncation parameter
    **kwargs : additional arguments
        Passed to compute_m_tests
        
    Returns
    -------
    float
        MZa statistic
    """
    results = compute_m_tests(X, trend=trend, detrend_method=detrend_method, k=k, **kwargs)
    return results["MZa"]


def MZt(
    X: np.ndarray,
    trend: str = "c",
    detrend_method: str = "ols",
    k: Optional[int] = None,
    **kwargs
) -> float:
    """
    Compute MZt unit root statistic.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend specification
    detrend_method : str, default="ols"
        Detrending method
    k : int, optional
        Lag truncation parameter
    **kwargs : additional arguments
        Passed to compute_m_tests
        
    Returns
    -------
    float
        MZt statistic
    """
    results = compute_m_tests(X, trend=trend, detrend_method=detrend_method, k=k, **kwargs)
    return results["MZt"]


def MSB(
    X: np.ndarray,
    trend: str = "c",
    detrend_method: str = "ols",
    k: Optional[int] = None,
    **kwargs
) -> float:
    """
    Compute MSB unit root statistic.
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend specification
    detrend_method : str, default="ols"
        Detrending method
    k : int, optional
        Lag truncation parameter
    **kwargs : additional arguments
        Passed to compute_m_tests
        
    Returns
    -------
    float
        MSB statistic
    """
    results = compute_m_tests(X, trend=trend, detrend_method=detrend_method, k=k, **kwargs)
    return results["MSB"]


def standard_m_test(
    X: np.ndarray,
    trend: str = "c",
    alpha: float = 0.05,
    **kwargs
) -> Dict[str, any]:
    """
    Perform standard M unit root test (assumes homoskedastic errors).
    
    Parameters
    ----------
    X : np.ndarray
        Time series data
    trend : str, default="c"
        Deterministic trend
    alpha : float, default=0.05
        Significance level
    **kwargs : additional arguments
        Passed to compute_m_tests
        
    Returns
    -------
    results : dict
        Test results including statistics, critical values, and decision
        
    Notes
    -----
    WARNING: This test is only valid under homoskedastic errors.
    Use bootstrap_test or simulation_test for non-stationary volatility.
    """
    from .critical_values import get_critical_value
    
    stats = compute_m_tests(X, trend=trend, **kwargs)
    
    # Get homoskedastic critical values
    cv_MZa = get_critical_value("MZa", trend=trend, alpha=alpha, variance_profile=None)
    cv_MZt = get_critical_value("MZt", trend=trend, alpha=alpha, variance_profile=None)
    cv_MSB = get_critical_value("MSB", trend=trend, alpha=alpha, variance_profile=None)
    
    # Test decisions (left-tailed for MZa and MZt, MSB rejects for small values)
    reject_MZa = stats["MZa"] < cv_MZa
    reject_MZt = stats["MZt"] < cv_MZt
    reject_MSB = stats["MSB"] < cv_MSB
    
    results = {
        "statistics": {
            "MZa": stats["MZa"],
            "MZt": stats["MZt"],
            "MSB": stats["MSB"],
        },
        "critical_values": {
            "MZa": cv_MZa,
            "MZt": cv_MZt,
            "MSB": cv_MSB,
        },
        "reject_null": {
            "MZa": reject_MZa,
            "MZt": reject_MZt,
            "MSB": reject_MSB,
        },
        "alpha": alpha,
        "trend": trend,
        "k": stats["k"],
        "warning": "Standard test assumes homoskedastic errors. "
                  "Use bootstrap or simulation-based tests for non-stationary volatility."
    }
    
    return results
