"""
Test ISTA algorithm implementation in GCONCORD.
"""

import pytest
import numpy as np
import gconcord as gc


class TestISTA:
    """Test ISTA algorithm functionality."""
    
    def setup_method(self):
        """Set up test data."""
        np.random.seed(456)
        self.p = 8
        A = np.random.randn(self.p, self.p)
        self.S = A @ A.T + 0.1 * np.eye(self.p)
        self.lambda_val = 0.1
        
    def test_ista_basic(self):
        """Test basic ISTA functionality."""
        result = gc.concord(self.S, lambda_param=self.lambda_val, algorithm='ista', max_iterations=1000, tolerance=5e-3)
        
        assert result['converged']
        assert result['algorithm_used'] == 'ista'
        assert result['omega'].shape == (self.p, self.p)
        
        # Check positive definiteness (all eigenvalues > 0)
        omega_dense = result['omega'].toarray()
        eigenvals = np.linalg.eigvals(omega_dense)
        assert np.all(eigenvals > 0)
        
    def test_ista_convergence_tracking(self):
        """Test ISTA with convergence tracking."""
        result = gc.concord(self.S, lambda_param=self.lambda_val, 
                          algorithm='ista', save_iterates=True, max_iterations=1000, tolerance=5e-3)
        
        assert result['objective_values'] is not None
        assert result['convergence_metrics'] is not None
        
        # Objective should generally decrease
        obj_vals = result['objective_values']
        if len(obj_vals) > 1:
            # ISTA can have non-monotonic convergence due to backtracking
            # Just check that final objective is less than initial
            assert obj_vals[-1] < obj_vals[0]
            
    def test_ista_iteration_limits(self):
        """Test ISTA with different iteration limits (since step size is not configurable)."""
        # Test with default iteration limit
        result_default = gc.concord(self.S, lambda_param=self.lambda_val, algorithm='ista', tolerance=5e-3)
        
        # Test with higher iteration limit
        result_more_iters = gc.concord(self.S, lambda_param=self.lambda_val, 
                                     algorithm='ista', max_iterations=1000, tolerance=5e-3)
        
        # Both should converge
        assert result_default['converged'] or result_more_iters['converged'], "At least one should converge"
        
        # Both should produce valid results if they converged
        if result_default['converged']:
            assert result_default['omega'] is not None
            assert result_default['algorithm_used'] == 'ista'
        if result_more_iters['converged']:
            assert result_more_iters['omega'] is not None
            assert result_more_iters['algorithm_used'] == 'ista'
        
    def test_ista_tolerance_settings(self):
        """Test ISTA with different tolerance settings."""
        # Loose tolerance
        result_loose = gc.concord(self.S, lambda_param=self.lambda_val, 
                                algorithm='ista', tolerance=1e-2, max_iterations=500)
        
        # Tight tolerance
        result_tight = gc.concord(self.S, lambda_param=self.lambda_val, 
                                algorithm='ista', tolerance=5e-3, max_iterations=1000)
        
        # At least one should converge
        assert result_loose['converged'] or result_tight['converged'], "At least one tolerance setting should work"
        
        # Loose tolerance should converge faster if both converge
        if result_loose['converged'] and result_tight['converged']:
            assert result_loose['iterations'] <= result_tight['iterations']
            
    def test_ista_matrix_lambda(self):
        """Test ISTA with matrix-valued lambda."""
        # Create custom penalty matrix
        Lambda = 0.1 * np.ones((self.p, self.p))
        np.fill_diagonal(Lambda, 0)
        
        # Make some entries zero (no penalty)
        Lambda[0, 1] = Lambda[1, 0] = 0
        Lambda[2, 3] = Lambda[3, 2] = 0
        
        # ISTA with matrix lambda can converge slower, use higher tolerance
        result = gc.concord(self.S, lambda_param=Lambda, algorithm='ista', max_iterations=1000, tolerance=1e-2)
        
        assert result['converged']
        omega = result['omega'].toarray()
        
        # Check that unpenalized entries are likely non-zero
        # (though not guaranteed due to other constraints)
        assert omega[0, 1] != 0 or omega[2, 3] != 0
        
    def test_ista_vs_ccista_consistency(self):
        """Test that ISTA through concord() matches ccista()."""
        # Run through concord interface
        result_concord = gc.concord(self.S, lambda_param=self.lambda_val, 
                                  algorithm='ista', tolerance=5e-3, max_iterations=1000)
        
        # Run through ccista interface
        omega_ccista, hist = gc.ccista(self.S, lambda1=self.lambda_val, 
                                      epstol=5e-3, maxitr=1000)
        
        # At least concord should converge
        assert result_concord['converged'], "ISTA through concord should converge"
        
        # Results should be similar (allowing for implementation differences)
        omega_concord = result_concord['omega'].toarray()
        omega_ccista_dense = omega_ccista.toarray()
        
        # Check that solutions are close
        np.testing.assert_allclose(omega_concord, omega_ccista_dense, rtol=1e-3, atol=1e-5)