"""
Critical values for unit root tests.

Provides critical values for M statistics under homoskedastic errors and
functions to generate critical values for heteroskedastic cases.
"""

import numpy as np
from typing import Dict, Optional, Callable

# Critical values from Elliott, Rothenberg & Stock (1996, Table 1)
# and Ng & Perron (2001, Table 1)
HOMOSKEDASTIC_CRITICAL_VALUES = {
    "c": {  # Constant only
        "MZa": {
            0.01: -13.8,
            0.05: -8.1,
            0.10: -5.7,
        },
        "MZt": {
            0.01: -2.58,
            0.05: -1.98,
            0.10: -1.62,
        },
        "MSB": {
            0.01: 0.174,
            0.05: 0.233,
            0.10: 0.275,
        },
    },
    "ct": {  # Constant and trend
        "MZa": {
            0.01: -23.8,
            0.05: -17.3,
            0.10: -14.2,
        },
        "MZt": {
            0.01: -3.42,
            0.05: -2.91,
            0.10: -2.62,
        },
        "MSB": {
            0.01: 0.143,
            0.05: 0.168,
            0.10: 0.185,
        },
    },
}


def get_critical_value(
    statistic: str,
    trend: str = "c",
    alpha: float = 0.05,
    variance_profile: Optional[Callable] = None,
) -> float:
    """
    Get critical value for M statistic.
    
    Parameters
    ----------
    statistic : str
        Test statistic: "MZa", "MZt", or "MSB"
    trend : str, default="c"
        Deterministic trend
    alpha : float, default=0.05
        Significance level
    variance_profile : callable, optional
        If provided, generates critical value for this variance profile
        via simulation. Otherwise returns homoskedastic critical value.
        
    Returns
    -------
    critical_value : float
        Critical value at specified significance level
    """
    if variance_profile is None:
        # Return homoskedastic critical value
        if alpha not in [0.01, 0.05, 0.10]:
            # Linear interpolation for other alpha values
            alphas = np.array([0.01, 0.05, 0.10])
            cvs = np.array([
                HOMOSKEDASTIC_CRITICAL_VALUES[trend][statistic][a] 
                for a in alphas
            ])
            critical_value = np.interp(alpha, alphas, cvs)
        else:
            critical_value = HOMOSKEDASTIC_CRITICAL_VALUES[trend][statistic][alpha]
    else:
        # Generate critical value for given variance profile
        from .simulation_tests import simulate_critical_values
        cv_dict = simulate_critical_values(
            variance_profile,
            statistic=statistic,
            trend=trend,
            alpha_levels=np.array([alpha]),
        )
        critical_value = cv_dict[alpha]
    
    return critical_value


def generate_critical_value_table(
    variance_profile: Callable,
    trend: str = "c",
    N: int = 50000,
    T: int = 1000,
    seed: Optional[int] = None,
) -> Dict[str, Dict[float, float]]:
    """
    Generate complete critical value table for given variance profile.
    
    Parameters
    ----------
    variance_profile : callable
        Variance profile function η(s)
    trend : str, default="c"
        Deterministic trend
    N : int, default=50000
        Number of simulations (use large N for accurate critical values)
    T : int, default=1000
        Grid size for Brownian motion
    seed : int, optional
        Random seed
        
    Returns
    -------
    critical_values : dict
        Nested dictionary: {statistic: {alpha: cv}}
        
    Notes
    -----
    Generates high-quality critical value table suitable for publication.
    Uses large N for accurate estimation.
    """
    from .simulation_tests import simulate_critical_values
    
    alpha_levels = np.array([0.01, 0.025, 0.05, 0.075, 0.10])
    statistics = ["MZa", "MZt", "MSB"]
    
    critical_values = {}
    
    for stat in statistics:
        cv_dict = simulate_critical_values(
            variance_profile,
            statistic=stat,
            trend=trend,
            N=N,
            T=T,
            alpha_levels=alpha_levels,
            seed=seed,
        )
        critical_values[stat] = cv_dict
    
    return critical_values
