"""
Test coordinate descent algorithms in GCONCORD.
"""

import pytest
import numpy as np
import gconcord as gc


class TestCoordinateDescent:
    """Test coordinate descent algorithm variants."""
    
    def setup_method(self):
        """Set up test data."""
        np.random.seed(123)
        # Standard case (n > p)
        self.p_standard = 10
        self.n_standard = 50
        A = np.random.randn(self.p_standard, self.p_standard)
        self.S_standard = A @ A.T + 0.1 * np.eye(self.p_standard)
        
        # High-dimensional case (n <= p)
        self.p_highdim = 20
        self.n_highdim = 10
        self.data_highdim = np.random.randn(self.n_highdim, self.p_highdim)
        self.S_highdim = np.cov(self.data_highdim.T)
        
    def test_coordinate_descent_standard(self):
        """Test standard coordinate descent algorithm."""
        result = gc.concord(self.S_standard, lambda_param=0.1, algorithm='coordinate', max_iterations=500)
        
        assert result['converged']
        assert result['algorithm_used'] in ['coordinate_descent', 'coordinate_blas']
        assert result['omega'].shape == (self.p_standard, self.p_standard)
        
        # Check symmetry
        omega_dense = result['omega'].toarray()
        np.testing.assert_allclose(omega_dense, omega_dense.T, rtol=1e-10)
        
    def test_coordinate_descent_blas(self):
        """Test BLAS-optimized coordinate descent."""
        result = gc.concord(self.S_standard, lambda_param=0.1, algorithm='coordinate_blas', max_iterations=500)
        
        assert result['converged']
        assert result['algorithm_used'] == 'coordinate_blas'
        
    def test_coordinate_descent_high_dim(self):
        """Test high-dimensional coordinate descent."""
        # This algorithm requires data matrix, not just covariance
        result = gc.concord(data=self.data_highdim, lambda_param=0.1, algorithm='coordinate_high_dim', max_iterations=500)
        
        assert result['omega'].shape == (self.p_highdim, self.p_highdim)
        
    def test_auto_algorithm_selection(self):
        """Test automatic algorithm selection."""
        # Small problem should select ISTA
        p_small = 5
        S_small = np.eye(p_small) + 0.1 * np.random.randn(p_small, p_small)
        S_small = (S_small + S_small.T) / 2
        
        result_small = gc.concord(S_small, lambda_param=0.1, algorithm='auto')
        assert result_small['algorithm_used'] in ['ista', 'coordinate_descent', 'coordinate_blas']
        
        # Large problem should select coordinate descent
        p_large = 100
        S_large = np.eye(p_large) + 0.01 * np.random.randn(p_large, p_large)
        S_large = (S_large + S_large.T) / 2
        
        result_large = gc.concord(S_large, lambda_param=0.1, algorithm='auto', max_iterations=10)
        assert result_large['algorithm_used'] in ['coordinate_descent', 'coordinate_blas']
        
    def test_lambda_scale_parameter(self):
        """Test lambda scaling parameter."""
        # Default lambda_scale = 2.0
        result_default = gc.concord(self.S_standard, lambda_param=0.1)
        
        # Custom lambda_scale
        result_custom = gc.concord(self.S_standard, lambda_param=0.1, lambda_scale=1.0)
        
        # Different scaling should produce different results
        omega_default = result_default['omega'].toarray()
        omega_custom = result_custom['omega'].toarray()
        
        assert not np.allclose(omega_default, omega_custom, rtol=1e-5)
        
    def test_penalize_diagonal_option(self):
        """Test penalize_diagonal option."""
        # Note: penalize_diagonal affects the Lambda matrix construction, 
        # but coordinate descent algorithms update diagonal elements independently
        # to maintain positive definiteness
        
        # Without diagonal penalization (default)
        result_no_pen = gc.concord(self.S_standard, lambda_param=0.1, penalize_diagonal=False, max_iterations=500)
        
        # With diagonal penalization
        result_with_pen = gc.concord(self.S_standard, lambda_param=0.1, penalize_diagonal=True, max_iterations=500)
        
        # Both should converge
        assert result_no_pen['converged']
        assert result_with_pen['converged']
        
        # Both should produce valid precision matrices
        assert result_no_pen['omega'] is not None
        assert result_with_pen['omega'] is not None