"""
Diagnostic Tests Module

This module provides diagnostic tests commonly used in conjunction with
asymmetric causality testing, including tests for multivariate normality
and ARCH effects.

Reference:
    Doornik, J.A. & Hansen, H. (2008). An Omnibus Test for Univariate and 
    Multivariate Normality. Oxford Bulletin of Economics and Statistics, 70, 927-939.
    
    Hacker, R.S. & Hatemi-J, A. (2005). A Multivariate Test for ARCH Effects.
    Applied Economics Letters, 12(7), 411-417.
"""

import numpy as np
from scipy import stats
from scipy.stats import chi2, jarque_bera
from typing import Tuple, Dict, Optional
from dataclasses import dataclass


@dataclass
class DiagnosticResult:
    """
    Data class for diagnostic test results.
    
    Attributes
    ----------
    test_name : str
        Name of the diagnostic test.
    statistic : float
        Test statistic value.
    p_value : float
        P-value of the test.
    df : int
        Degrees of freedom (if applicable).
    reject_null : bool
        Whether null hypothesis is rejected at 5% level.
    details : Dict
        Additional test-specific details.
    """
    test_name: str
    statistic: float
    p_value: float
    df: int
    reject_null: bool
    details: Dict
    
    def __repr__(self) -> str:
        return (f"DiagnosticResult(test='{self.test_name}', "
                f"statistic={self.statistic:.4f}, p_value={self.p_value:.4f})")
    
    def summary(self) -> str:
        """Generate formatted summary."""
        conclusion = "Rejected" if self.reject_null else "Not Rejected"
        return (f"{self.test_name}\n"
                f"  Test Statistic: {self.statistic:.4f}\n"
                f"  P-value: {self.p_value:.4f}\n"
                f"  Degrees of Freedom: {self.df}\n"
                f"  H0 at 5%: {conclusion}")


def doornik_hansen_test(residuals: np.ndarray) -> DiagnosticResult:
    """
    Doornik-Hansen test for multivariate normality.
    
    This test is an omnibus test for univariate and multivariate normality
    based on transformed skewness and kurtosis.
    
    Parameters
    ----------
    residuals : np.ndarray
        T x K matrix of residuals from VAR model.
    
    Returns
    -------
    DiagnosticResult
        Test results including statistic, p-value, and conclusion.
    
    Notes
    -----
    H0: The residuals are multivariate normal.
    H1: The residuals are not multivariate normal.
    
    Reference:
        Doornik, J.A. & Hansen, H. (2008). An Omnibus Test for Univariate 
        and Multivariate Normality. Oxford Bulletin of Economics and 
        Statistics, 70, 927-939.
    """
    residuals = np.asarray(residuals)
    if residuals.ndim == 1:
        residuals = residuals.reshape(-1, 1)
    
    T, K = residuals.shape
    
    # Standardize residuals
    mean_res = np.mean(residuals, axis=0)
    std_res = np.std(residuals, axis=0, ddof=1)
    z = (residuals - mean_res) / std_res
    
    # Calculate skewness and kurtosis for each variable
    skew = np.mean(z**3, axis=0)
    kurt = np.mean(z**4, axis=0)
    
    # Transformation parameters (Doornik-Hansen, 2008)
    n = T
    beta = 3 * (n**2 + 27*n - 70) * (n + 1) * (n + 3) / ((n - 2) * (n + 5) * (n + 7) * (n + 9))
    w2 = -1 + np.sqrt(2 * (beta - 1))
    delta = 1 / np.sqrt(np.log(np.sqrt(w2)))
    alpha = np.sqrt(2 / (w2 - 1))
    
    # Transform skewness
    y = skew * np.sqrt((n + 1) * (n + 3) / (6 * (n - 2)))
    z1 = delta * np.log(y / alpha + np.sqrt((y / alpha)**2 + 1))
    
    # Transform kurtosis
    delta_k = (n - 3) * (n + 1) * (n**2 + 15*n - 4)
    a = ((n - 2) * (n + 5) * (n + 7) * (n**2 + 27*n - 70)) / (6 * delta_k)
    c = ((n - 7) * (n + 5) * (n + 7) * (n**2 + 2*n - 5)) / (6 * delta_k)
    k = ((n + 5) * (n + 7) * (n**3 + 37*n**2 + 11*n - 313)) / (12 * delta_k)
    alpha_k = a + skew**2 * c
    
    chi = (kurt - 1 - skew**2) * 2 * k
    chi = np.maximum(chi, 1e-10)  # Ensure positive
    
    z2 = (((chi / (2 * alpha_k))**(1/3)) - 1 + 1 / (9 * alpha_k)) * np.sqrt(9 * alpha_k)
    
    # Test statistic
    Ep = np.sum(z1**2 + z2**2)
    df = 2 * K
    p_value = 1 - chi2.cdf(Ep, df)
    
    return DiagnosticResult(
        test_name="Doornik-Hansen Multivariate Normality Test",
        statistic=Ep,
        p_value=p_value,
        df=df,
        reject_null=p_value < 0.05,
        details={'skewness_transform': z1.tolist(), 
                 'kurtosis_transform': z2.tolist()}
    )


def multivariate_jarque_bera(residuals: np.ndarray) -> DiagnosticResult:
    """
    Multivariate Jarque-Bera test for normality.
    
    A multivariate extension of the Jarque-Bera test based on the sum
    of univariate JB statistics.
    
    Parameters
    ----------
    residuals : np.ndarray
        T x K matrix of residuals.
    
    Returns
    -------
    DiagnosticResult
        Test results.
    """
    residuals = np.asarray(residuals)
    if residuals.ndim == 1:
        residuals = residuals.reshape(-1, 1)
    
    T, K = residuals.shape
    
    # Calculate JB statistic for each variable
    jb_stats = []
    for k in range(K):
        _, p = jarque_bera(residuals[:, k])
        s = stats.skew(residuals[:, k])
        ku = stats.kurtosis(residuals[:, k], fisher=True)
        jb = (T / 6) * (s**2 + (ku**2) / 4)
        jb_stats.append(jb)
    
    # Sum of JB statistics
    total_jb = np.sum(jb_stats)
    df = 2 * K
    p_value = 1 - chi2.cdf(total_jb, df)
    
    return DiagnosticResult(
        test_name="Multivariate Jarque-Bera Normality Test",
        statistic=total_jb,
        p_value=p_value,
        df=df,
        reject_null=p_value < 0.05,
        details={'univariate_jb': jb_stats}
    )


def multivariate_arch_test(residuals: np.ndarray, 
                           lags: int = 5,
                           n_bootstrap: int = 1000,
                           random_state: Optional[int] = None) -> DiagnosticResult:
    """
    Multivariate ARCH test with bootstrap critical values.
    
    This test is based on Hacker and Hatemi-J (2005) and tests for
    ARCH effects in VAR residuals.
    
    Parameters
    ----------
    residuals : np.ndarray
        T x K matrix of VAR residuals.
    lags : int, optional
        Number of lags for ARCH test. Default is 5.
    n_bootstrap : int, optional
        Number of bootstrap replications. Default is 1000.
    random_state : int, optional
        Random seed for reproducibility.
    
    Returns
    -------
    DiagnosticResult
        Test results with bootstrap p-value.
    
    Notes
    -----
    H0: No ARCH effects (homoskedasticity).
    H1: ARCH effects present (heteroskedasticity).
    
    Reference:
        Hacker, R.S. & Hatemi-J, A. (2005). A Multivariate Test for 
        ARCH Effects. Applied Economics Letters, 12(7), 411-417.
    """
    if random_state is not None:
        np.random.seed(random_state)
    
    residuals = np.asarray(residuals)
    if residuals.ndim == 1:
        residuals = residuals.reshape(-1, 1)
    
    T, K = residuals.shape
    
    # Calculate squared residuals
    squared_res = residuals**2
    
    # Vectorize squared residuals (vech operation for each time point)
    # For simplicity, use full vectorization
    vech_squared = squared_res  # K columns
    
    # Calculate test statistic using Lagrange Multiplier approach
    def calc_lm_stat(sq_res):
        T_eff = sq_res.shape[0] - lags
        if T_eff <= 0:
            return 0
        
        # Create lagged variables
        Y = sq_res[lags:, :]
        X = np.ones((T_eff, 1))
        
        for lag in range(1, lags + 1):
            X = np.column_stack([X, sq_res[lags - lag:T - lag, :]])
        
        # Regress squared residuals on lagged squared residuals
        try:
            beta = np.linalg.lstsq(X, Y, rcond=None)[0]
            fitted = X @ beta
            ss_res = np.sum((Y - fitted)**2)
            ss_tot = np.sum((Y - Y.mean(axis=0))**2)
            
            if ss_tot < 1e-10:
                return 0
            
            r_squared = 1 - ss_res / ss_tot
            lm_stat = T_eff * K * r_squared
            return lm_stat
        except:
            return 0
    
    # Calculate test statistic for original data
    lm_original = calc_lm_stat(squared_res)
    
    # Bootstrap
    boot_stats = np.zeros(n_bootstrap)
    centered_res = residuals - residuals.mean(axis=0)
    
    for b in range(n_bootstrap):
        # Resample residuals
        indices = np.random.randint(0, T, size=T)
        boot_res = centered_res[indices, :]
        boot_squared = boot_res**2
        boot_stats[b] = calc_lm_stat(boot_squared)
    
    # Calculate bootstrap p-value
    p_value = np.mean(boot_stats >= lm_original)
    
    # Asymptotic df
    df = K * K * lags
    asymp_p = 1 - chi2.cdf(lm_original, df)
    
    return DiagnosticResult(
        test_name="Multivariate ARCH Test (Hacker-Hatemi-J)",
        statistic=lm_original,
        p_value=p_value,
        df=df,
        reject_null=p_value < 0.05,
        details={'bootstrap_p_value': p_value, 
                 'asymptotic_p_value': asymp_p,
                 'n_bootstrap': n_bootstrap,
                 'lags': lags}
    )


def ljung_box_test(residuals: np.ndarray, 
                   lags: int = 10) -> DiagnosticResult:
    """
    Multivariate Ljung-Box test for autocorrelation.
    
    Parameters
    ----------
    residuals : np.ndarray
        T x K matrix of residuals.
    lags : int, optional
        Number of lags to test. Default is 10.
    
    Returns
    -------
    DiagnosticResult
        Test results.
    """
    residuals = np.asarray(residuals)
    if residuals.ndim == 1:
        residuals = residuals.reshape(-1, 1)
    
    T, K = residuals.shape
    
    # Calculate sample autocorrelations
    Q = 0
    for h in range(1, lags + 1):
        # Cross-correlation matrix at lag h
        Gamma_h = (residuals[h:, :].T @ residuals[:-h, :]) / T
        Gamma_0 = (residuals.T @ residuals) / T
        
        try:
            Gamma_0_inv = np.linalg.inv(Gamma_0)
            C_h = Gamma_h @ Gamma_0_inv @ Gamma_h.T @ Gamma_0_inv
            Q += np.trace(C_h) / (T - h)
        except:
            pass
    
    Q = T * (T + 2) * Q
    df = K * K * lags
    p_value = 1 - chi2.cdf(Q, df)
    
    return DiagnosticResult(
        test_name="Multivariate Ljung-Box Test",
        statistic=Q,
        p_value=p_value,
        df=df,
        reject_null=p_value < 0.05,
        details={'lags': lags}
    )


def run_all_diagnostics(residuals: np.ndarray,
                        arch_lags: int = 5,
                        lb_lags: int = 10,
                        n_bootstrap: int = 1000,
                        random_state: Optional[int] = None) -> Dict[str, DiagnosticResult]:
    """
    Run all diagnostic tests on VAR residuals.
    
    Parameters
    ----------
    residuals : np.ndarray
        T x K matrix of VAR residuals.
    arch_lags : int, optional
        Number of lags for ARCH test. Default is 5.
    lb_lags : int, optional
        Number of lags for Ljung-Box test. Default is 10.
    n_bootstrap : int, optional
        Bootstrap replications for ARCH test. Default is 1000.
    random_state : int, optional
        Random seed.
    
    Returns
    -------
    Dict[str, DiagnosticResult]
        Dictionary with results for all diagnostic tests.
    """
    results = {}
    
    # Normality tests
    results['doornik_hansen'] = doornik_hansen_test(residuals)
    results['jarque_bera'] = multivariate_jarque_bera(residuals)
    
    # ARCH test
    results['arch'] = multivariate_arch_test(
        residuals, lags=arch_lags, 
        n_bootstrap=n_bootstrap, 
        random_state=random_state
    )
    
    # Autocorrelation test
    results['ljung_box'] = ljung_box_test(residuals, lags=lb_lags)
    
    return results


def print_diagnostics_summary(results: Dict[str, DiagnosticResult]) -> str:
    """
    Generate formatted summary of all diagnostic tests.
    
    Parameters
    ----------
    results : Dict[str, DiagnosticResult]
        Dictionary of diagnostic test results.
    
    Returns
    -------
    str
        Formatted summary string.
    """
    separator = "=" * 70
    output = f"""
{separator}
                    DIAGNOSTIC TEST RESULTS
{separator}

{'Test':<40}{'Statistic':>12}{'p-value':>12}{'H0':>10}
{'-' * 74}
"""
    
    for name, result in results.items():
        h0_status = "Reject" if result.reject_null else "Accept"
        output += f"{result.test_name[:38]:<40}{result.statistic:>12.4f}"
        output += f"{result.p_value:>12.4f}{h0_status:>10}\n"
    
    output += f"""
{'-' * 74}
Note: H0 rejection based on 5% significance level.
{separator}
"""
    
    return output
