"""
Test suite for diagnostic and unit root tests.
"""

import numpy as np
import pytest
from numpy.testing import assert_array_almost_equal

from asymcaus import (
    # Diagnostic tests
    doornik_hansen_test,
    multivariate_jarque_bera,
    multivariate_arch_test,
    ljung_box_test,
    run_all_diagnostics,
    DiagnosticResult,
    
    # Unit root tests
    adf_test,
    kpss_test,
    pp_test,
    determine_integration_order,
    UnitRootResult,
)


class TestNormalityTests:
    """Tests for normality testing functions."""
    
    def test_doornik_hansen_normal_data(self):
        """Test Doornik-Hansen with normal data."""
        np.random.seed(42)
        residuals = np.random.randn(200, 2)
        
        result = doornik_hansen_test(residuals)
        
        assert isinstance(result, DiagnosticResult)
        assert result.test_name == "Doornik-Hansen Multivariate Normality Test"
        # For truly normal data, p-value should be relatively high
        # (though random chance may cause rejection)
        assert 0 <= result.p_value <= 1
    
    def test_doornik_hansen_non_normal_data(self):
        """Test Doornik-Hansen with non-normal data."""
        np.random.seed(42)
        # Generate heavily skewed data
        residuals = np.random.exponential(1, (200, 2)) - 1
        
        result = doornik_hansen_test(residuals)
        
        assert isinstance(result, DiagnosticResult)
        # Should likely reject normality for exponential data
        # (p-value should be small)
    
    def test_jarque_bera_normal_data(self):
        """Test multivariate Jarque-Bera with normal data."""
        np.random.seed(42)
        residuals = np.random.randn(200, 2)
        
        result = multivariate_jarque_bera(residuals)
        
        assert isinstance(result, DiagnosticResult)
        assert result.test_name == "Multivariate Jarque-Bera Normality Test"
        assert np.isfinite(result.statistic)
    
    def test_1d_input(self):
        """Test with 1D input."""
        np.random.seed(42)
        residuals = np.random.randn(100)
        
        result = doornik_hansen_test(residuals)
        assert isinstance(result, DiagnosticResult)


class TestARCHTest:
    """Tests for ARCH effect testing."""
    
    def test_arch_no_effects(self):
        """Test ARCH test on homoskedastic data."""
        np.random.seed(42)
        residuals = np.random.randn(200, 2)
        
        result = multivariate_arch_test(
            residuals, 
            lags=3, 
            n_bootstrap=100,
            random_state=42
        )
        
        assert isinstance(result, DiagnosticResult)
        assert 'ARCH' in result.test_name
        # Should likely not reject for homoskedastic data
    
    def test_arch_with_effects(self):
        """Test ARCH test on heteroskedastic data."""
        np.random.seed(42)
        T = 200
        
        # Generate ARCH(1) process
        residuals = np.zeros((T, 2))
        sigma2 = np.ones((T, 2))
        
        for t in range(1, T):
            sigma2[t] = 0.1 + 0.8 * residuals[t-1]**2
            residuals[t] = np.sqrt(sigma2[t]) * np.random.randn(2)
        
        result = multivariate_arch_test(
            residuals,
            lags=3,
            n_bootstrap=100,
            random_state=42
        )
        
        assert isinstance(result, DiagnosticResult)


class TestLjungBoxTest:
    """Tests for Ljung-Box autocorrelation test."""
    
    def test_no_autocorrelation(self):
        """Test with white noise (no autocorrelation)."""
        np.random.seed(42)
        residuals = np.random.randn(200, 2)
        
        result = ljung_box_test(residuals, lags=5)
        
        assert isinstance(result, DiagnosticResult)
        assert 'Ljung-Box' in result.test_name
    
    def test_with_autocorrelation(self):
        """Test with autocorrelated residuals."""
        np.random.seed(42)
        T = 200
        
        # Generate AR(1) process
        residuals = np.zeros((T, 2))
        for t in range(1, T):
            residuals[t] = 0.7 * residuals[t-1] + np.random.randn(2) * 0.5
        
        result = ljung_box_test(residuals, lags=5)
        
        assert isinstance(result, DiagnosticResult)


class TestRunAllDiagnostics:
    """Tests for comprehensive diagnostics."""
    
    def test_all_diagnostics(self):
        """Test running all diagnostics."""
        np.random.seed(42)
        residuals = np.random.randn(200, 2)
        
        results = run_all_diagnostics(
            residuals,
            arch_lags=3,
            lb_lags=5,
            n_bootstrap=50,
            random_state=42
        )
        
        assert isinstance(results, dict)
        assert 'doornik_hansen' in results
        assert 'jarque_bera' in results
        assert 'arch' in results
        assert 'ljung_box' in results
        
        for name, result in results.items():
            assert isinstance(result, DiagnosticResult)


class TestADFTest:
    """Tests for Augmented Dickey-Fuller test."""
    
    def test_stationary_series(self):
        """Test ADF with stationary series."""
        np.random.seed(42)
        y = np.random.randn(200)  # White noise is stationary
        
        result = adf_test(y, trend='c')
        
        assert isinstance(result, UnitRootResult)
        assert result.test_name == "Augmented Dickey-Fuller Test"
        # Should reject unit root for stationary series
        assert result.statistic < result.critical_values['5%']
    
    def test_unit_root_series(self):
        """Test ADF with unit root series."""
        np.random.seed(42)
        y = np.cumsum(np.random.randn(200))  # Random walk has unit root
        
        result = adf_test(y, trend='c')
        
        assert isinstance(result, UnitRootResult)
        # May or may not reject - depends on specific realization
    
    def test_different_trends(self):
        """Test ADF with different trend specifications."""
        np.random.seed(42)
        y = np.random.randn(200)
        
        for trend in ['n', 'c', 'ct']:
            result = adf_test(y, trend=trend)
            assert isinstance(result, UnitRootResult)
            assert np.isfinite(result.statistic)


class TestKPSSTest:
    """Tests for KPSS stationarity test."""
    
    def test_stationary_series(self):
        """Test KPSS with stationary series."""
        np.random.seed(42)
        y = np.random.randn(200)
        
        result = kpss_test(y, trend='c')
        
        assert isinstance(result, UnitRootResult)
        assert result.test_name == "KPSS Test"
        # Should NOT reject stationarity for white noise
    
    def test_non_stationary_series(self):
        """Test KPSS with non-stationary series."""
        np.random.seed(42)
        y = np.cumsum(np.random.randn(200))
        
        result = kpss_test(y, trend='c')
        
        assert isinstance(result, UnitRootResult)


class TestPPTest:
    """Tests for Phillips-Perron test."""
    
    def test_stationary_series(self):
        """Test PP with stationary series."""
        np.random.seed(42)
        y = np.random.randn(200)
        
        result = pp_test(y, trend='c')
        
        assert isinstance(result, UnitRootResult)
        assert result.test_name == "Phillips-Perron Test"
    
    def test_unit_root_series(self):
        """Test PP with unit root series."""
        np.random.seed(42)
        y = np.cumsum(np.random.randn(200))
        
        result = pp_test(y, trend='c')
        
        assert isinstance(result, UnitRootResult)


class TestIntegrationOrder:
    """Tests for integration order determination."""
    
    def test_i0_series(self):
        """Test I(0) series detection."""
        np.random.seed(42)
        y = np.random.randn(200)  # Stationary
        
        order, results = determine_integration_order(y)
        
        assert order == 0
        assert 'd0' in results
    
    def test_i1_series(self):
        """Test I(1) series detection."""
        np.random.seed(42)
        y = np.cumsum(np.random.randn(200))  # First difference is stationary
        
        order, results = determine_integration_order(y)
        
        # Should be 1 (or possibly 0 if random walk happens to look stationary)
        assert order in [0, 1]
        assert 'd0' in results
    
    def test_max_order_limit(self):
        """Test maximum order limit."""
        np.random.seed(42)
        y = np.cumsum(np.cumsum(np.random.randn(200)))  # I(2)
        
        order, results = determine_integration_order(y, max_order=1)
        
        assert order <= 1


class TestUnitRootResultSummary:
    """Tests for UnitRootResult summary method."""
    
    def test_summary_generation(self):
        """Test summary string generation."""
        np.random.seed(42)
        y = np.random.randn(100)
        
        result = adf_test(y)
        summary = result.summary()
        
        assert isinstance(summary, str)
        assert 'Augmented Dickey-Fuller' in summary
        assert 'Critical Values' in summary
        assert 'Conclusion' in summary


class TestDiagnosticResultSummary:
    """Tests for DiagnosticResult summary method."""
    
    def test_summary_generation(self):
        """Test summary string generation."""
        np.random.seed(42)
        residuals = np.random.randn(100, 2)
        
        result = doornik_hansen_test(residuals)
        summary = result.summary()
        
        assert isinstance(summary, str)
        assert 'Doornik-Hansen' in summary


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
