"""
Test Suite for Shor's Algorithm Implementation
===============================================

Comprehensive tests for quantum period finding, RSA factorization,
resource estimation, and success probability calculations.
"""

import pytest
from datetime import datetime
from chronos.core.simulate.shors_algorithm import (
    ShorsAlgorithm,
    QuantumPeriodFinder,
    ClassicalPostProcessor,
    QuantumResourceEstimator,
    QuantumCircuitVisualizer,
    SuccessProbabilityCalculator,
    AlgorithmPhase,
    RSAKeySize,
    QubitQuality,
    FactorMethod,
    create_shors_algorithm,
)


class TestQuantumPeriodFinder:
    """Test quantum period finding functionality."""
    
    def test_period_finder_initialization(self):
        """Test PeriodFinder can be initialized."""
        finder = QuantumPeriodFinder()
        assert finder is not None
    
    def test_gcd_calculation(self):
        """Test GCD calculation."""
        finder = QuantumPeriodFinder()
        assert finder._gcd(12, 8) == 4
        assert finder._gcd(17, 5) == 1
        assert finder._gcd(100, 50) == 50
    
    def test_modular_exponentiation(self):
        """Test fast modular exponentiation."""
        finder = QuantumPeriodFinder()
        # 2^3 mod 5 = 8 mod 5 = 3
        assert finder._modular_exponentiation(2, 3, 5) == 3
        # 3^4 mod 7 = 81 mod 7 = 4
        assert finder._modular_exponentiation(3, 4, 7) == 4
    
    def test_find_order_small_numbers(self):
        """Test finding order with small numbers."""
        finder = QuantumPeriodFinder()
        # Order of 2 modulo 7
        order = finder.find_order(2, 7)
        assert order > 0
        # Verify: 2^order ≡ 1 (mod 7)
        assert finder._modular_exponentiation(2, order, 7) == 1
    
    def test_quantum_order_finding_success(self):
        """Test quantum order finding succeeds."""
        finder = QuantumPeriodFinder()
        result = finder.quantum_order_finding(2, 15, QubitQuality.PERFECT)
        assert result.n == 15
        assert result.a == 2
        assert result.period > 0
    
    def test_quantum_order_finding_not_coprime(self):
        """Test handling of non-coprime a and n."""
        finder = QuantumPeriodFinder()
        result = finder.quantum_order_finding(5, 15, QubitQuality.PERFECT)
        # gcd(5, 15) = 5, so should fail
        assert result.success == False


class TestClassicalPostProcessor:
    """Test classical post-processing functionality."""
    
    def test_post_processor_initialization(self):
        """Test PostProcessor can be initialized."""
        processor = ClassicalPostProcessor()
        assert processor is not None
    
    def test_gcd_calculation(self):
        """Test GCD in post-processor."""
        processor = ClassicalPostProcessor()
        assert processor._gcd(12, 8) == 4
        assert processor._gcd(21, 14) == 7
    
    def test_modular_exponentiation(self):
        """Test modular exponentiation in post-processor."""
        processor = ClassicalPostProcessor()
        assert processor._modular_exponentiation(2, 10, 1000) == 24
    
    def test_factor_extraction_from_period(self):
        """Test factor extraction given a period."""
        processor = ClassicalPostProcessor()
        # For n = 15, a = 4, period should be 2
        # 4^2 = 16 ≡ 1 (mod 15)
        p, q, success = processor.extract_factors_from_period(15, 4, 2)
        if success:
            assert p * q == 15
    
    def test_pollard_rho_small_number(self):
        """Test Pollard's rho algorithm."""
        processor = ClassicalPostProcessor()
        p, q, success, iterations = processor.pollard_rho(15)
        if success:
            assert p * q == 15
            assert p > 1 and q > 1
    
    def test_pollard_rho_even_number(self):
        """Test Pollard's rho with even number."""
        processor = ClassicalPostProcessor()
        p, q, success, iterations = processor.pollard_rho(10)
        assert success == True
        assert p == 2
        assert q == 5
    
    def test_verify_factorization(self):
        """Test factorization verification."""
        processor = ClassicalPostProcessor()
        assert processor.verify_factorization(15, 3, 5) == True
        assert processor.verify_factorization(15, 3, 4) == False
        assert processor.verify_factorization(15, 15, 1) == False


class TestShorsAlgorithm:
    """Test Shor's algorithm implementation."""
    
    def test_shors_initialization(self):
        """Test Shor's algorithm initialization."""
        algorithm = ShorsAlgorithm()
        assert algorithm is not None
        assert algorithm.algorithm_id is not None
    
    def test_shors_with_custom_id(self):
        """Test Shor's algorithm with custom ID."""
        algorithm = ShorsAlgorithm(algorithm_id="test-algo-1")
        assert algorithm.algorithm_id == "test-algo-1"
    
    def test_factor_even_number(self):
        """Test factoring even number (trivial case)."""
        algorithm = ShorsAlgorithm()
        result = algorithm.factor_rsa_key(10)
        assert result.success == True
        assert result.p == 2
        assert result.q == 5
    
    def test_factor_semiprime(self):
        """Test factoring semiprime."""
        algorithm = ShorsAlgorithm()
        # 15 = 3 * 5
        result = algorithm.factor_rsa_key(15, max_attempts=3)
        if result.success:
            assert result.p * result.q == 15
            assert result.p > 1 and result.q > 1
    
    def test_factor_semiprime_21(self):
        """Test factoring 21 = 3 * 7."""
        algorithm = ShorsAlgorithm()
        result = algorithm.factor_rsa_key(21, max_attempts=5)
        if result.success:
            assert result.p * result.q == 21
    
    def test_execution_log_created(self):
        """Test execution log is created."""
        algorithm = ShorsAlgorithm()
        algorithm.factor_rsa_key(15)
        log = algorithm.get_execution_log()
        assert log is not None
        assert log.n == 15
        assert log.start_time is not None
    
    def test_qubit_quality_affects_success(self):
        """Test that qubit quality affects results."""
        # Perfect qubits should succeed better than poor
        algorithm_perfect = ShorsAlgorithm()
        algorithm_poor = ShorsAlgorithm()
        
        result_perfect = algorithm_perfect.factor_rsa_key(21, max_attempts=3, qubit_quality=QubitQuality.PERFECT)
        result_poor = algorithm_poor.factor_rsa_key(21, max_attempts=3, qubit_quality=QubitQuality.POOR)
        
        # Perfect should have better or equal probability
        assert result_perfect.probability >= result_poor.probability


class TestResourceEstimation:
    """Test resource estimation functionality."""
    
    def test_circuit_resources_512bit(self):
        """Test circuit resource estimation for 512-bit."""
        circuit = QuantumResourceEstimator.estimate_circuit_resources(512)
        assert circuit.total_qubits > 0
        assert circuit.control_qubits == 2 * 512
        assert circuit.target_qubits == 512
        assert circuit.ancilla_qubits == 2 * 512 + 3
        assert circuit.depth > 0
        assert circuit.gates_count > 0
    
    def test_circuit_resources_768bit(self):
        """Test circuit resource estimation for 768-bit."""
        circuit = QuantumResourceEstimator.estimate_circuit_resources(768)
        assert circuit.total_qubits > 0
        assert circuit.control_qubits == 2 * 768
        assert circuit.depth > 0
    
    def test_resource_estimate_perfect_qubits(self):
        """Test resource estimate with perfect qubits."""
        estimate = QuantumResourceEstimator.estimate_resources(512, QubitQuality.PERFECT)
        assert estimate.key_size == 512
        assert estimate.logical_qubits_needed > 0
        assert estimate.physical_qubits_needed > 0
        assert estimate.error_correction_overhead == 1
        assert estimate.success_probability > 0.9
    
    def test_resource_estimate_poor_qubits(self):
        """Test resource estimate with poor qubits."""
        estimate = QuantumResourceEstimator.estimate_resources(512, QubitQuality.POOR)
        assert estimate.error_correction_overhead > 1
        assert estimate.success_probability < 0.5
    
    def test_error_correction_overhead_scaling(self):
        """Test error correction overhead scales with quality."""
        perfect = QuantumResourceEstimator.estimate_resources(512, QubitQuality.PERFECT)
        poor = QuantumResourceEstimator.estimate_resources(512, QubitQuality.POOR)
        assert poor.error_correction_overhead > perfect.error_correction_overhead
        assert poor.physical_qubits_needed > perfect.physical_qubits_needed
    
    def test_resource_estimates_different_key_sizes(self):
        """Test resource estimates for different key sizes."""
        size_512 = QuantumResourceEstimator.estimate_resources(512, QubitQuality.GOOD)
        size_768 = QuantumResourceEstimator.estimate_resources(768, QubitQuality.GOOD)
        size_2048 = QuantumResourceEstimator.estimate_resources(2048, QubitQuality.GOOD)
        
        assert size_512.logical_qubits_needed < size_768.logical_qubits_needed
        assert size_768.logical_qubits_needed < size_2048.logical_qubits_needed


class TestQuantumCircuitVisualizer:
    """Test quantum circuit visualization."""
    
    def test_draw_circuit_ascii(self):
        """Test ASCII circuit drawing."""
        circuit = QuantumCircuitVisualizer.draw_circuit_ascii(8)
        assert circuit is not None
        assert len(circuit) > 0
        assert "SHOR'S ALGORITHM" in circuit
        assert "Qubits" in circuit
    
    def test_draw_algorithm_flow(self):
        """Test algorithm flow drawing."""
        flow = QuantumCircuitVisualizer.draw_algorithm_flow()
        assert flow is not None
        assert len(flow) > 0
        assert "Shor's Algorithm" in flow or "SHOR'S" in flow
    
    def test_visualization_completeness(self):
        """Test visualizations contain key information."""
        circuit = QuantumCircuitVisualizer.draw_circuit_ascii(16)
        flow = QuantumCircuitVisualizer.draw_algorithm_flow()
        
        # Circuit should have gate information
        assert "Gates" in circuit
        
        # Flow should have algorithm steps
        assert "INPUT" in flow or "input" in flow


class TestSuccessProbabilityCalculator:
    """Test success probability calculations."""
    
    def test_probability_perfect_qubits(self):
        """Test probability with perfect qubits."""
        prob = SuccessProbabilityCalculator.calculate_probability(
            512, QubitQuality.PERFECT, num_attempts=1
        )
        assert 0.9 <= prob <= 1.0
    
    def test_probability_poor_qubits(self):
        """Test probability with poor qubits."""
        prob = SuccessProbabilityCalculator.calculate_probability(
            512, QubitQuality.POOR, num_attempts=1
        )
        assert 0.0 <= prob <= 0.5
    
    def test_probability_increases_with_attempts(self):
        """Test probability increases with more attempts."""
        prob_1 = SuccessProbabilityCalculator.calculate_probability(
            512, QubitQuality.GOOD, num_attempts=1
        )
        prob_5 = SuccessProbabilityCalculator.calculate_probability(
            512, QubitQuality.GOOD, num_attempts=5
        )
        prob_10 = SuccessProbabilityCalculator.calculate_probability(
            512, QubitQuality.GOOD, num_attempts=10
        )
        
        assert prob_1 < prob_5 < prob_10
    
    def test_expected_attempts_calculation(self):
        """Test expected attempts calculation."""
        attempts = SuccessProbabilityCalculator.calculate_expected_attempts(0.95)
        assert attempts > 0
        assert attempts < 10
    
    def test_probability_table(self):
        """Test probability table generation."""
        table = SuccessProbabilityCalculator.get_probability_table(512)
        assert len(table) == 6  # 6 quality levels
        assert all(isinstance(p, float) for p in table.values())
        assert all(0 <= p <= 1 for p in table.values())


class TestEnumerations:
    """Test enumeration definitions."""
    
    def test_algorithm_phases(self):
        """Test algorithm phase enumeration."""
        assert AlgorithmPhase.INITIALIZATION.value == "initialization"
        assert AlgorithmPhase.QUANTUM_PERIOD_FINDING.value == "quantum_period_finding"
        assert AlgorithmPhase.COMPLETED.value == "completed"
    
    def test_rsa_key_sizes(self):
        """Test RSA key size enumeration."""
        assert RSAKeySize.RSA_512.value == 512
        assert RSAKeySize.RSA_768.value == 768
        assert RSAKeySize.RSA_2048.value == 2048
    
    def test_qubit_quality_levels(self):
        """Test qubit quality enumeration."""
        qualities = [q for q in QubitQuality]
        assert len(qualities) == 6
        assert QubitQuality.PERFECT in qualities
        assert QubitQuality.POOR in qualities
    
    def test_factor_methods(self):
        """Test factorization method enumeration."""
        methods = [m for m in FactorMethod]
        assert len(methods) == 4
        assert FactorMethod.CONTINUED_FRACTIONS in methods
        assert FactorMethod.POLLARD_RHO in methods


class TestDataClasses:
    """Test data class functionality."""
    
    def test_quantum_circuit_info_creation(self):
        """Test QuantumCircuitInfo creation."""
        circuit = QuantumCircuitVisualizer.draw_circuit_ascii(8)
        assert circuit is not None
    
    def test_resource_estimate_creation(self):
        """Test ResourceEstimate creation."""
        estimate = QuantumResourceEstimator.estimate_resources(512, QubitQuality.GOOD)
        assert estimate.key_size == 512
        assert estimate.created_at is not None
    
    def test_factorization_creation(self):
        """Test Factorization result creation."""
        algorithm = ShorsAlgorithm()
        result = algorithm.factor_rsa_key(15)
        assert result.n == 15
        assert result.created_at is not None


class TestIntegration:
    """Integration tests for Shor's algorithm."""
    
    def test_factory_function(self):
        """Test factory function creates algorithm."""
        algorithm = create_shors_algorithm()
        assert algorithm is not None
        assert isinstance(algorithm, ShorsAlgorithm)
    
    def test_complete_workflow_15(self):
        """Test complete workflow for factoring 15."""
        algorithm = ShorsAlgorithm("test-15")
        result = algorithm.factor_rsa_key(15, max_attempts=5)
        log = algorithm.get_execution_log()
        
        assert log is not None
        assert log.algorithm_id == "test-15"
        assert log.n == 15
        assert log.end_time is not None
    
    def test_complete_workflow_21(self):
        """Test complete workflow for factoring 21."""
        algorithm = ShorsAlgorithm("test-21")
        result = algorithm.factor_rsa_key(21, max_attempts=5)
        log = algorithm.get_execution_log()
        
        assert log.n == 21
    
    def test_algorithm_phases_progression(self):
        """Test algorithm phases progress correctly."""
        algorithm = ShorsAlgorithm()
        algorithm.factor_rsa_key(15)
        log = algorithm.get_execution_log()
        
        assert AlgorithmPhase.INITIALIZATION in log.phases_completed or len(log.phases_completed) > 0


class TestEdgeCases:
    """Test edge cases and error handling."""
    
    def test_factor_large_prime(self):
        """Test that large prime cannot be factored easily."""
        algorithm = ShorsAlgorithm()
        prime = 29  # Prime number
        result = algorithm.factor_rsa_key(prime, max_attempts=1)
        # Should either fail or take special handling
        assert result is not None
    
    def test_factor_small_semiprime(self):
        """Test factoring small semiprimes."""
        algorithm = ShorsAlgorithm()
        for semiprime in [6, 10, 14, 15, 21, 22]:
            result = algorithm.factor_rsa_key(semiprime, max_attempts=3)
            assert result is not None
    
    def test_thread_safety(self):
        """Test thread safety of algorithm."""
        algorithm = ShorsAlgorithm()
        # Should be safe to call factor_rsa_key multiple times
        result1 = algorithm.factor_rsa_key(15)
        log1 = algorithm.get_execution_log()
        
        result2 = algorithm.factor_rsa_key(21)
        log2 = algorithm.get_execution_log()
        
        assert log2.n == 21  # Should have latest log


# Run tests if executed directly
if __name__ == "__main__":
    pytest.main([__file__, "-v"])
