"""
Test suite for asymcaus package.

Tests the core functionality including:
- Cumulative component calculation
- VAR estimation
- Lag selection
- Wald test
- Bootstrap critical values
- Full asymmetric causality test
"""

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

from asymcaus import (
    asymmetric_causality_test,
    full_asymmetric_causality_analysis,
    cumulative_components,
    var_lags,
    estimate_var,
    select_lag_order,
    information_criterion,
    wald_test,
    create_restriction_matrix,
    AsymmetricCausalityResult,
)


class TestCumulativeComponents:
    """Tests for cumulative_components function."""
    
    def test_basic_decomposition(self):
        """Test basic positive/negative decomposition."""
        y = np.array([1.0, 2.0, 1.5, 3.0, 2.5])
        pos, neg = cumulative_components(y)
        
        # First differences: [1.0, -0.5, 1.5, -0.5]
        # Positive cumsum: [1.0, 1.0, 2.5, 2.5]
        # Negative cumsum: [0.0, -0.5, -0.5, -1.0]
        
        expected_pos = np.array([[1.0], [1.0], [2.5], [2.5]])
        expected_neg = np.array([[0.0], [-0.5], [-0.5], [-1.0]])
        
        assert_array_almost_equal(pos, expected_pos)
        assert_array_almost_equal(neg, expected_neg)
    
    def test_2d_input(self):
        """Test with 2D input array."""
        y = np.array([[1.0, 10.0],
                      [2.0, 12.0],
                      [1.5, 11.0],
                      [3.0, 15.0]])
        
        pos, neg = cumulative_components(y)
        
        assert pos.shape == (3, 2)
        assert neg.shape == (3, 2)
    
    def test_all_positive_changes(self):
        """Test series with only positive changes."""
        y = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
        pos, neg = cumulative_components(y)
        
        # All changes are +1
        expected_pos = np.array([[1], [2], [3], [4]])
        expected_neg = np.zeros((4, 1))
        
        assert_array_almost_equal(pos, expected_pos)
        assert_array_almost_equal(neg, expected_neg)
    
    def test_all_negative_changes(self):
        """Test series with only negative changes."""
        y = np.array([5.0, 4.0, 3.0, 2.0, 1.0])
        pos, neg = cumulative_components(y)
        
        expected_pos = np.zeros((4, 1))
        expected_neg = np.array([[-1], [-2], [-3], [-4]])
        
        assert_array_almost_equal(pos, expected_pos)
        assert_array_almost_equal(neg, expected_neg)
    
    def test_log_transform(self):
        """Test with log transformation."""
        y = np.array([1.0, 2.0, 4.0, 8.0])
        pos, neg = cumulative_components(y, log_transform=True)
        
        # Log changes should be ~0.693 each
        assert pos.shape == (3, 1)
        assert np.all(pos > 0)


class TestVarLags:
    """Tests for var_lags function."""
    
    def test_single_lag(self):
        """Test with single lag."""
        data = np.array([[1, 10],
                         [2, 20],
                         [3, 30],
                         [4, 40],
                         [5, 50]])
        
        y, y_lags = var_lags(data, lags=1)
        
        assert y.shape == (4, 2)
        assert y_lags.shape == (4, 2)
        
        # First row of y should be [2, 20]
        assert_array_almost_equal(y[0], [2, 20])
        # First row of lags should be [1, 10]
        assert_array_almost_equal(y_lags[0], [1, 10])
    
    def test_multiple_lags(self):
        """Test with multiple lags."""
        data = np.array([[1, 10],
                         [2, 20],
                         [3, 30],
                         [4, 40],
                         [5, 50]])
        
        y, y_lags = var_lags(data, lags=2)
        
        assert y.shape == (3, 2)
        assert y_lags.shape == (3, 4)  # 2 vars × 2 lags = 4 columns


class TestLagSelection:
    """Tests for lag selection functions."""
    
    def test_select_lag_order(self):
        """Test optimal lag selection."""
        np.random.seed(42)
        T = 200
        
        # Generate VAR(2) process
        data = np.zeros((T, 2))
        data[0:2] = np.random.randn(2, 2)
        
        for t in range(2, T):
            data[t] = 0.3 * data[t-1] + 0.2 * data[t-2] + np.random.randn(2) * 0.5
        
        optimal_lag, _ = select_lag_order(data, max_lag=5, criterion='bic')
        
        # Should select lag close to true order (2)
        assert 1 <= optimal_lag <= 4
    
    def test_information_criteria(self):
        """Test different information criteria."""
        varcov = np.eye(2) * 0.5
        T = 100
        n_vars = 2
        lag = 2
        
        aic = information_criterion(varcov, T, n_vars, lag, 'aic')
        bic = information_criterion(varcov, T, n_vars, lag, 'bic')
        hqc = information_criterion(varcov, T, n_vars, lag, 'hqc')
        hjc = information_criterion(varcov, T, n_vars, lag, 'hjc')
        
        # All should be finite
        assert np.isfinite(aic)
        assert np.isfinite(bic)
        assert np.isfinite(hqc)
        assert np.isfinite(hjc)
        
        # HJC should be average of BIC and HQC
        assert_allclose(hjc, (bic + hqc) / 2, rtol=1e-10)


class TestVAREstimation:
    """Tests for VAR estimation."""
    
    def test_unrestricted_estimation(self):
        """Test unrestricted VAR estimation."""
        np.random.seed(42)
        T = 100
        y = np.random.randn(T, 2)
        X = np.column_stack([np.ones(T), np.random.randn(T, 4)])
        
        Ahat, leverage = estimate_var(y, X, restrict=False)
        
        assert Ahat.shape == (2, 5)  # 2 equations, 5 params each
        
        # Check predictions are reasonable
        fitted = X @ Ahat.T
        assert fitted.shape == y.shape
    
    def test_restricted_estimation(self):
        """Test restricted VAR estimation."""
        np.random.seed(42)
        T = 100
        y = np.random.randn(T, 2)
        X = np.column_stack([np.ones(T), np.random.randn(T, 4)])
        
        rvector = np.array([0, 0, 1, 0, 1])  # Restrict 3rd and 5th coefficients
        
        Ahat, leverage = estimate_var(y, X, restrict=True, rvector=rvector)
        
        assert Ahat.shape == (2, 5)
        # Restricted coefficients should be zero
        assert Ahat[0, 2] == 0
        assert Ahat[0, 4] == 0


class TestRestrictionMatrix:
    """Tests for restriction matrix creation."""
    
    def test_basic_restrictions(self):
        """Test basic restriction matrix."""
        rvector, rmatrix = create_restriction_matrix(n_vars=2, var_order=2, add_lags=0)
        
        # Should have restrictions for 2 lags
        assert rmatrix.shape[0] == 2
    
    def test_with_additional_lags(self):
        """Test with Toda-Yamamoto additional lags."""
        rvector, rmatrix = create_restriction_matrix(n_vars=2, var_order=2, add_lags=1)
        
        # rvector should cover more parameters
        assert len(rvector) == 1 + 2 * 3  # constant + 2 vars × 3 total lags


class TestWaldTest:
    """Tests for Wald test."""
    
    def test_wald_statistic(self):
        """Test Wald statistic calculation."""
        np.random.seed(42)
        T = 100
        y = np.random.randn(T, 2)
        X = np.column_stack([np.ones(T), np.random.randn(T, 2)])
        
        Ahat = np.linalg.lstsq(X, y, rcond=None)[0].T
        rmatrix = np.array([[0, 0, 0, 0, 0, 1]])
        
        wald_stat = wald_test(y, X, Ahat, rmatrix)
        
        assert np.isfinite(wald_stat)
        assert wald_stat >= 0


class TestAsymmetricCausalityTest:
    """Tests for the main asymmetric causality test function."""
    
    @pytest.fixture
    def sample_data(self):
        """Generate sample data for testing."""
        np.random.seed(42)
        T = 300
        z = np.cumsum(np.random.randn(T))
        y = np.cumsum(np.random.randn(T))
        return y, z
    
    def test_positive_component(self, sample_data):
        """Test positive component causality."""
        y, z = sample_data
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            n_bootstrap=100,
            random_state=42
        )
        
        assert isinstance(result, AsymmetricCausalityResult)
        assert np.isfinite(result.wald_stat)
        assert 0 <= result.p_value_chi2 <= 1
        assert len(result.bootstrap_cv) == 3
        assert result.optimal_lag >= 1
    
    def test_negative_component(self, sample_data):
        """Test negative component causality."""
        y, z = sample_data
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='negative',
            n_bootstrap=100,
            random_state=42
        )
        
        assert isinstance(result, AsymmetricCausalityResult)
        assert 'negative' in result.component_type.lower()
    
    def test_original_variables(self, sample_data):
        """Test causality with original variables."""
        y, z = sample_data
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='original',
            n_bootstrap=100,
            random_state=42
        )
        
        assert isinstance(result, AsymmetricCausalityResult)
    
    def test_with_integration_order(self, sample_data):
        """Test with Toda-Yamamoto procedure."""
        y, z = sample_data
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            integration_order=1,
            n_bootstrap=100,
            random_state=42
        )
        
        assert isinstance(result, AsymmetricCausalityResult)
    
    def test_variable_names(self, sample_data):
        """Test custom variable names."""
        y, z = sample_data
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            var_names=('Oil', 'Stock'),
            n_bootstrap=100,
            random_state=42
        )
        
        assert 'Oil' in result.variable_names or 'Stock' in result.variable_names
    
    def test_different_criteria(self, sample_data):
        """Test different information criteria."""
        y, z = sample_data
        
        for criterion in ['aic', 'bic', 'hqc', 'hjc']:
            result = asymmetric_causality_test(
                y=y, z=z,
                component='positive',
                criterion=criterion,
                n_bootstrap=50,
                random_state=42
            )
            assert isinstance(result, AsymmetricCausalityResult)
    
    def test_summary_output(self, sample_data):
        """Test summary string generation."""
        y, z = sample_data
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            n_bootstrap=100,
            random_state=42
        )
        
        summary = result.summary()
        assert isinstance(summary, str)
        assert 'Wald' in summary
        assert 'Bootstrap' in summary


class TestFullAnalysis:
    """Tests for full asymmetric causality analysis."""
    
    def test_full_analysis(self):
        """Test comprehensive analysis."""
        np.random.seed(42)
        T = 200
        z = np.cumsum(np.random.randn(T))
        y = np.cumsum(np.random.randn(T))
        
        results = full_asymmetric_causality_analysis(
            y=y, z=z,
            n_bootstrap=50,
            random_state=42
        )
        
        assert isinstance(results, dict)
        assert 'positive' in results
        assert 'negative' in results
        assert 'original' in results
        assert 'positive_to_negative' in results
        assert 'negative_to_positive' in results
        
        for key, result in results.items():
            assert isinstance(result, AsymmetricCausalityResult)


class TestEdgeCases:
    """Tests for edge cases and error handling."""
    
    def test_short_series(self):
        """Test with short time series."""
        np.random.seed(42)
        y = np.cumsum(np.random.randn(50))
        z = np.cumsum(np.random.randn(50))
        
        result = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            max_lag=3,
            n_bootstrap=50,
            random_state=42
        )
        
        assert isinstance(result, AsymmetricCausalityResult)
    
    def test_unequal_lengths_error(self):
        """Test error for unequal length series."""
        y = np.random.randn(100)
        z = np.random.randn(90)
        
        with pytest.raises(ValueError):
            asymmetric_causality_test(y=y, z=z)
    
    def test_invalid_component_error(self):
        """Test error for invalid component type."""
        y = np.random.randn(100)
        z = np.random.randn(100)
        
        with pytest.raises(ValueError):
            asymmetric_causality_test(y=y, z=z, component='invalid')


class TestReproducibility:
    """Tests for reproducibility with random state."""
    
    def test_same_seed_same_results(self):
        """Test that same seed produces same results."""
        np.random.seed(42)
        y = np.cumsum(np.random.randn(200))
        z = np.cumsum(np.random.randn(200))
        
        result1 = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            n_bootstrap=100,
            random_state=123
        )
        
        result2 = asymmetric_causality_test(
            y=y, z=z,
            component='positive',
            n_bootstrap=100,
            random_state=123
        )
        
        assert result1.wald_stat == result2.wald_stat
        assert_array_almost_equal(result1.bootstrap_cv, result2.bootstrap_cv)


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