"""
CHRONOS Shor's Algorithm for RSA Factorization
===============================================

Complete implementation of Shor's algorithm for quantum period finding
and RSA key factorization. Includes classical post-processing, success
probability estimation, and quantum resource requirements.
"""

from dataclasses import dataclass
from datetime import datetime
from enum import Enum
import math
from typing import Optional, Tuple, List, Dict, Any
import hashlib
from random import randint, choice
from threading import RLock


# ============================================================================
# Enumerations
# ============================================================================

class AlgorithmPhase(Enum):
    """Shor's algorithm execution phases."""
    INITIALIZATION = "initialization"
    QUANTUM_PERIOD_FINDING = "quantum_period_finding"
    CLASSICAL_PROCESSING = "classical_processing"
    FACTOR_VERIFICATION = "factor_verification"
    COMPLETED = "completed"
    FAILED = "failed"


class RSAKeySize(Enum):
    """Standard RSA key sizes."""
    RSA_512 = 512
    RSA_768 = 768
    RSA_1024 = 1024
    RSA_2048 = 2048
    RSA_4096 = 4096


class QubitQuality(Enum):
    """Qubit quality/error rate classifications."""
    PERFECT = "perfect"  # 0% error
    EXCELLENT = "excellent"  # 0.01% error per gate
    VERY_GOOD = "very_good"  # 0.1% error per gate
    GOOD = "good"  # 0.5% error per gate
    FAIR = "fair"  # 1% error per gate
    POOR = "poor"  # 5% error per gate


class FactorMethod(Enum):
    """Classical post-processing methods."""
    CONTINUED_FRACTIONS = "continued_fractions"
    POLLARD_RHO = "pollard_rho"
    TRIAL_DIVISION = "trial_division"
    QUADRATIC_SIEVE = "quadratic_sieve"


# ============================================================================
# Data Classes
# ============================================================================

@dataclass
class QuantumCircuitInfo:
    """Information about quantum circuit requirements."""
    total_qubits: int
    control_qubits: int
    target_qubits: int
    ancilla_qubits: int
    depth: int
    gates_count: int
    cnot_count: int
    single_qubit_gates: int
    circuit_name: str = "Shor's Algorithm Circuit"
    created_at: datetime = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()


@dataclass
class ResourceEstimate:
    """Resource estimation for Shor's algorithm execution."""
    key_size: int
    logical_qubits_needed: int
    physical_qubits_needed: int  # Accounting for error correction
    circuit_depth: int
    gate_count: int
    cnot_gates: int
    single_qubit_gates: int
    estimated_runtime_seconds: float
    estimated_error_rate: float
    success_probability: float
    qubit_quality: QubitQuality
    error_correction_overhead: int
    created_at: datetime = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()


@dataclass
class Factorization:
    """Result of RSA factorization."""
    n: int  # The number to factor
    p: int  # First prime factor
    q: int  # Second prime factor
    success: bool
    method_used: FactorMethod
    iterations: int
    execution_time_ms: float
    attempts: int
    probability: float
    verification_passed: bool
    created_at: datetime = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()


@dataclass
class PeriodFindingResult:
    """Result of quantum period finding."""
    n: int  # The number N
    a: int  # Random base
    period: int  # Found period r
    success: bool
    measurement_result: int
    denominator: int
    numerator: int
    iterations: int
    amplitude_amplification_rounds: int
    created_at: datetime = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()


@dataclass
class ShorsExecutionLog:
    """Execution log for Shor's algorithm."""
    algorithm_id: str
    start_time: datetime
    end_time: Optional[datetime]
    n: int  # Number to factor
    key_size: int
    current_phase: AlgorithmPhase
    phases_completed: List[AlgorithmPhase]
    period_finding_results: List[PeriodFindingResult]
    factorization_result: Optional[Factorization]
    total_iterations: int
    success: bool
    error_message: Optional[str]


# ============================================================================
# Quantum Period Finding (Classical Simulation)
# ============================================================================

class QuantumPeriodFinder:
    """Simulates quantum period finding using order-finding algorithm."""
    
    def __init__(self):
        self.lock = RLock()
    
    def _gcd(self, a: int, b: int) -> int:
        """Calculate greatest common divisor."""
        while b:
            a, b = b, a % b
        return a
    
    def _modular_exponentiation(self, base: int, exp: int, mod: int) -> int:
        """Efficient modular exponentiation using binary exponentiation."""
        result = 1
        base = base % mod
        while exp > 0:
            if exp % 2 == 1:
                result = (result * base) % mod
            exp = exp >> 1
            base = (base * base) % mod
        return result
    
    def find_order(self, a: int, n: int) -> int:
        """
        Find the order (period) of a modulo n.
        Order r is the smallest positive integer where a^r ≡ 1 (mod n).
        """
        if self._gcd(a, n) != 1:
            return -1  # a is not coprime with n
        
        r = 1
        result = a
        max_iterations = 2 * n  # Safety limit
        
        for _ in range(max_iterations):
            result = (result * a) % n
            if result == 1:
                return r
            r += 1
        
        return -1  # Order not found
    
    def quantum_order_finding(self, a: int, n: int, qubit_quality: QubitQuality = QubitQuality.GOOD) -> PeriodFindingResult:
        """
        Simulate quantum period finding with error modeling based on qubit quality.
        
        This is a classical simulation of the quantum Fourier transform approach.
        In a real quantum computer, this would use quantum phase estimation.
        """
        with self.lock:
            # Find the actual order
            r = self.find_order(a, n)
            
            if r == -1:
                return PeriodFindingResult(
                    n=n,
                    a=a,
                    period=0,
                    success=False,
                    measurement_result=0,
                    denominator=0,
                    numerator=0,
                    iterations=2*n,
                    amplitude_amplification_rounds=0
                )
            
            # Simulate quantum measurement with error based on qubit quality
            error_rate = self._get_error_rate_for_quality(qubit_quality)
            
            # Quantum measurement process (simulated)
            measurement_result = self._simulate_quantum_measurement(r, error_rate)
            
            # Extract fraction using continued fractions algorithm
            numerator, denominator = self._continued_fractions(measurement_result, 2**int(math.log2(n)))
            
            # Determine if we got the correct period
            success = (denominator == r) or (denominator * 2 == r)
            
            return PeriodFindingResult(
                n=n,
                a=a,
                period=r,
                success=success,
                measurement_result=measurement_result,
                denominator=denominator,
                numerator=numerator,
                iterations=int(math.log2(n)) + 1,
                amplitude_amplification_rounds=int(math.log2(n))
            )
    
    def _get_error_rate_for_quality(self, quality: QubitQuality) -> float:
        """Get error rate for given qubit quality."""
        error_rates = {
            QubitQuality.PERFECT: 0.0,
            QubitQuality.EXCELLENT: 0.0001,
            QubitQuality.VERY_GOOD: 0.001,
            QubitQuality.GOOD: 0.005,
            QubitQuality.FAIR: 0.01,
            QubitQuality.POOR: 0.05,
        }
        return error_rates.get(quality, 0.01)
    
    def _simulate_quantum_measurement(self, period: int, error_rate: float) -> int:
        """Simulate quantum measurement with error."""
        if error_rate == 0:
            return period
        
        # Simulate measurement error
        import random
        if random.random() < error_rate:
            # Return incorrect measurement
            return max(1, period + random.randint(-period // 2, period // 2))
        
        return period
    
    def _continued_fractions(self, numerator: int, denominator: int, max_depth: int = 20) -> Tuple[int, int]:
        """
        Extract period from fraction using continued fractions algorithm.
        Used in classical post-processing of quantum measurement.
        """
        cf = []
        a, b = numerator, denominator
        
        for _ in range(max_depth):
            q = a // b
            cf.append(q)
            a, b = b, a - q * b
            if b == 0:
                break
        
        # Compute convergents
        if len(cf) == 0:
            return 0, 1
        
        h_prev, h_curr = 1, cf[0]
        k_prev, k_curr = 0, 1
        
        for i in range(1, len(cf)):
            h_prev, h_curr = h_curr, cf[i] * h_curr + h_prev
            k_prev, k_curr = k_curr, cf[i] * k_curr + k_prev
        
        return h_curr, k_curr


# ============================================================================
# Classical Post-Processing
# ============================================================================

class ClassicalPostProcessor:
    """Classical post-processing for Shor's algorithm."""
    
    def __init__(self):
        self.lock = RLock()
    
    def _gcd(self, a: int, b: int) -> int:
        """Calculate GCD."""
        while b:
            a, b = b, a % b
        return a
    
    def _modular_exponentiation(self, base: int, exp: int, mod: int) -> int:
        """Fast modular exponentiation."""
        result = 1
        base = base % mod
        while exp > 0:
            if exp % 2 == 1:
                result = (result * base) % mod
            exp = exp >> 1
            base = (base * base) % mod
        return result
    
    def extract_factors_from_period(self, n: int, a: int, r: int) -> Tuple[int, int, bool]:
        """
        Classical factor extraction given period r.
        If a^r ≡ 1 (mod n), then factors can often be found.
        """
        if r % 2 != 0:
            return 0, 0, False
        
        # Compute a^(r/2) mod n
        x = self._modular_exponentiation(a, r // 2, n)
        
        if x == -1 % n or x == 1:
            return 0, 0, False
        
        # Try to factor using GCD
        factor1 = self._gcd(x - 1, n)
        factor2 = self._gcd(x + 1, n)
        
        if 1 < factor1 < n:
            return factor1, n // factor1, True
        
        if 1 < factor2 < n:
            return factor2, n // factor2, True
        
        return 0, 0, False
    
    def pollard_rho(self, n: int, max_iterations: int = 100000) -> Tuple[int, int, bool, int]:
        """
        Pollard's rho algorithm for integer factorization.
        Used as fallback when period-based factoring fails.
        """
        if n % 2 == 0:
            return 2, n // 2, True, 1
        
        def f(x: int) -> int:
            return (x * x + 1) % n
        
        x = randint(2, n - 1)
        y = x
        c = randint(1, n - 1)
        d = 1
        iterations = 0
        
        while d == 1 and iterations < max_iterations:
            x = f(x)
            y = f(f(y))
            d = self._gcd(abs(x - y), n)
            iterations += 1
        
        if d == n:
            return 0, 0, False, iterations
        
        return d, n // d, True, iterations
    
    def continued_fractions_factoring(self, n: int, max_iterations: int = 1000) -> Tuple[int, int, bool, int]:
        """
        Attempt factoring using convergents of continued fraction expansion.
        """
        # Try to find factors using continued fractions
        for b in range(2, int(math.sqrt(n)) + 1):
            cf = self._get_continued_fraction_convergents(n, b, max_iterations)
            for num, den in cf:
                d = self._gcd(den, n)
                if 1 < d < n:
                    return d, n // d, True, b
        
        return 0, 0, False, 0
    
    def _get_continued_fraction_convergents(self, n: int, k: int, max_depth: int) -> List[Tuple[int, int]]:
        """Get convergents of continued fraction."""
        convergents = []
        
        # Simplified continued fraction calculation
        a, b = k, 1
        for _ in range(max_depth):
            if b == 0:
                break
            convergents.append((a, b))
            a, b = b, a - (a // b) * b
        
        return convergents
    
    def verify_factorization(self, n: int, p: int, q: int) -> bool:
        """Verify that p * q == n."""
        return p * q == n and p > 1 and q > 1 and p != q


# ============================================================================
# Shor's Algorithm Implementation
# ============================================================================

class ShorsAlgorithm:
    """Complete implementation of Shor's algorithm for RSA factorization."""
    
    def __init__(self, algorithm_id: str = None):
        self.algorithm_id = algorithm_id or hashlib.sha256(str(datetime.now()).encode()).hexdigest()[:8]
        self.period_finder = QuantumPeriodFinder()
        self.post_processor = ClassicalPostProcessor()
        self.lock = RLock()
        self.execution_log = None
    
    def _gcd(self, a: int, b: int) -> int:
        """Calculate GCD."""
        while b:
            a, b = b, a % b
        return a
    
    def _is_prime_power(self, n: int) -> bool:
        """Check if n is a perfect prime power."""
        for p in range(2, int(math.sqrt(n)) + 1):
            if n % p == 0:
                # Check if n = p^k
                temp = n
                while temp % p == 0:
                    temp //= p
                return temp == 1
        return False
    
    def factor_rsa_key(self, n: int, max_attempts: int = 5, 
                      qubit_quality: QubitQuality = QubitQuality.GOOD) -> Factorization:
        """
        Factor RSA modulus n using Shor's algorithm.
        
        Args:
            n: RSA modulus to factor
            max_attempts: Maximum attempts (uses different random bases)
            qubit_quality: Simulated qubit quality
            
        Returns:
            Factorization result
        """
        with self.lock:
            start_time = datetime.now()
            
            # Initialize execution log
            self.execution_log = ShorsExecutionLog(
                algorithm_id=self.algorithm_id,
                start_time=start_time,
                end_time=None,
                n=n,
                key_size=n.bit_length(),
                current_phase=AlgorithmPhase.INITIALIZATION,
                phases_completed=[],
                period_finding_results=[],
                factorization_result=None,
                total_iterations=0,
                success=False,
                error_message=None
            )
            
            # Check if n is even
            if n % 2 == 0:
                end_time = datetime.now()
                result = Factorization(
                    n=n,
                    p=2,
                    q=n // 2,
                    success=True,
                    method_used=FactorMethod.TRIAL_DIVISION,
                    iterations=1,
                    execution_time_ms=(end_time - start_time).total_seconds() * 1000,
                    attempts=1,
                    probability=1.0,
                    verification_passed=True
                )
                self.execution_log.factorization_result = result
                self.execution_log.success = True
                self.execution_log.end_time = end_time
                return result
            
            # Check if n is a perfect power
            if self._is_prime_power(n):
                self.execution_log.error_message = "N is a perfect prime power"
                self.execution_log.current_phase = AlgorithmPhase.FAILED
                self.execution_log.end_time = datetime.now()
                return Factorization(
                    n=n, p=0, q=0, success=False,
                    method_used=FactorMethod.TRIAL_DIVISION,
                    iterations=0, execution_time_ms=0,
                    attempts=1, probability=0.0,
                    verification_passed=False
                )
            
            # Main factoring loop
            for attempt in range(max_attempts):
                self.execution_log.current_phase = AlgorithmPhase.QUANTUM_PERIOD_FINDING
                
                # Choose random base
                a = randint(2, n - 1)
                
                # Check if gcd(a, n) != 1
                d = self._gcd(a, n)
                if d != 1 and d != n:
                    result = Factorization(
                        n=n, p=d, q=n // d,
                        success=True,
                        method_used=FactorMethod.TRIAL_DIVISION,
                        iterations=1,
                        execution_time_ms=(datetime.now() - start_time).total_seconds() * 1000,
                        attempts=attempt + 1,
                        probability=1.0,
                        verification_passed=True
                    )
                    self.execution_log.factorization_result = result
                    self.execution_log.success = True
                    self.execution_log.current_phase = AlgorithmPhase.COMPLETED
                    self.execution_log.end_time = datetime.now()
                    return result
                
                # Quantum period finding
                period_result = self.period_finder.quantum_order_finding(a, n, qubit_quality)
                self.execution_log.period_finding_results.append(period_result)
                
                if not period_result.success:
                    continue
                
                # Classical post-processing
                self.execution_log.current_phase = AlgorithmPhase.CLASSICAL_PROCESSING
                
                # Try to extract factors from period
                p, q, success = self.post_processor.extract_factors_from_period(n, a, period_result.period)
                
                if success:
                    self.execution_log.current_phase = AlgorithmPhase.FACTOR_VERIFICATION
                    
                    if self.post_processor.verify_factorization(n, p, q):
                        end_time = datetime.now()
                        result = Factorization(
                            n=n, p=p, q=q,
                            success=True,
                            method_used=FactorMethod.CONTINUED_FRACTIONS,
                            iterations=attempt + 1,
                            execution_time_ms=(end_time - start_time).total_seconds() * 1000,
                            attempts=attempt + 1,
                            probability=self._calculate_success_probability(qubit_quality),
                            verification_passed=True
                        )
                        self.execution_log.factorization_result = result
                        self.execution_log.success = True
                        self.execution_log.current_phase = AlgorithmPhase.COMPLETED
                        self.execution_log.end_time = end_time
                        return result
            
            # Fallback: Try Pollard's rho if period finding didn't work
            self.execution_log.current_phase = AlgorithmPhase.CLASSICAL_PROCESSING
            p, q, success, iterations = self.post_processor.pollard_rho(n)
            
            if success and self.post_processor.verify_factorization(n, p, q):
                end_time = datetime.now()
                result = Factorization(
                    n=n, p=p, q=q,
                    success=True,
                    method_used=FactorMethod.POLLARD_RHO,
                    iterations=max_attempts + iterations,
                    execution_time_ms=(end_time - start_time).total_seconds() * 1000,
                    attempts=max_attempts + 1,
                    probability=0.5,
                    verification_passed=True
                )
                self.execution_log.factorization_result = result
                self.execution_log.success = True
                self.execution_log.current_phase = AlgorithmPhase.COMPLETED
                self.execution_log.end_time = end_time
                return result
            
            # Factoring failed
            end_time = datetime.now()
            self.execution_log.current_phase = AlgorithmPhase.FAILED
            self.execution_log.error_message = f"Failed to factor after {max_attempts} attempts"
            self.execution_log.end_time = end_time
            
            return Factorization(
                n=n, p=0, q=0,
                success=False,
                method_used=FactorMethod.CONTINUED_FRACTIONS,
                iterations=max_attempts,
                execution_time_ms=(end_time - start_time).total_seconds() * 1000,
                attempts=max_attempts,
                probability=0.0,
                verification_passed=False
            )
    
    def _calculate_success_probability(self, qubit_quality: QubitQuality) -> float:
        """Calculate success probability based on qubit quality."""
        probabilities = {
            QubitQuality.PERFECT: 0.99,
            QubitQuality.EXCELLENT: 0.95,
            QubitQuality.VERY_GOOD: 0.85,
            QubitQuality.GOOD: 0.70,
            QubitQuality.FAIR: 0.50,
            QubitQuality.POOR: 0.20,
        }
        return probabilities.get(qubit_quality, 0.5)
    
    def get_execution_log(self) -> ShorsExecutionLog:
        """Get execution log of last algorithm run."""
        return self.execution_log


# ============================================================================
# Resource Estimation
# ============================================================================

class QuantumResourceEstimator:
    """Estimate quantum resources needed for Shor's algorithm."""
    
    @staticmethod
    def estimate_circuit_resources(key_size: int) -> QuantumCircuitInfo:
        """
        Estimate quantum circuit resources for given RSA key size.
        Based on Shor's algorithm resource requirements.
        """
        # Control qubits (for superposition): 2n
        control_qubits = 2 * key_size
        
        # Target qubits (for register): n
        target_qubits = key_size
        
        # Ancilla qubits (for auxiliary operations): 2n + 3
        ancilla_qubits = 2 * key_size + 3
        
        # Total logical qubits
        total_qubits = control_qubits + target_qubits + ancilla_qubits
        
        # Circuit depth (number of time steps): O(n^3)
        depth = key_size ** 3
        
        # Gate count: O(n^3)
        gates_count = key_size ** 3
        
        # CNOT gates: ~70% of gates
        cnot_count = int(gates_count * 0.7)
        
        # Single qubit gates
        single_qubit_gates = gates_count - cnot_count
        
        return QuantumCircuitInfo(
            total_qubits=total_qubits,
            control_qubits=control_qubits,
            target_qubits=target_qubits,
            ancilla_qubits=ancilla_qubits,
            depth=depth,
            gates_count=gates_count,
            cnot_count=cnot_count,
            single_qubit_gates=single_qubit_gates,
            circuit_name=f"Shor's Algorithm for RSA-{key_size}"
        )
    
    @staticmethod
    def estimate_resources(key_size: int, qubit_quality: QubitQuality) -> ResourceEstimate:
        """
        Estimate all resources needed for Shor's algorithm execution.
        """
        circuit_info = QuantumResourceEstimator.estimate_circuit_resources(key_size)
        
        # Error correction overhead based on qubit quality
        error_rates = {
            QubitQuality.PERFECT: 1,
            QubitQuality.EXCELLENT: 10,
            QubitQuality.VERY_GOOD: 100,
            QubitQuality.GOOD: 1000,
            QubitQuality.FAIR: 5000,
            QubitQuality.POOR: 50000,
        }
        error_correction_overhead = error_rates.get(qubit_quality, 1000)
        
        # Physical qubits needed (logical qubits * error correction overhead)
        physical_qubits = circuit_info.total_qubits * error_correction_overhead
        
        # Estimate error rate per gate
        error_rate_per_gate = {
            QubitQuality.PERFECT: 0.0,
            QubitQuality.EXCELLENT: 0.0001,
            QubitQuality.VERY_GOOD: 0.001,
            QubitQuality.GOOD: 0.01,
            QubitQuality.FAIR: 0.05,
            QubitQuality.POOR: 0.1,
        }.get(qubit_quality, 0.01)
        
        # Overall circuit error rate
        circuit_error_rate = 1 - (1 - error_rate_per_gate) ** circuit_info.gates_count
        
        # Estimate runtime (seconds): depth / gate_time (assume 100ns per gate)
        gate_time_ns = 100
        estimated_runtime = (circuit_info.depth * gate_time_ns) / 1e9
        
        # Success probability
        success_prob = {
            QubitQuality.PERFECT: 0.99,
            QubitQuality.EXCELLENT: 0.95,
            QubitQuality.VERY_GOOD: 0.85,
            QubitQuality.GOOD: 0.70,
            QubitQuality.FAIR: 0.50,
            QubitQuality.POOR: 0.20,
        }.get(qubit_quality, 0.5)
        
        return ResourceEstimate(
            key_size=key_size,
            logical_qubits_needed=circuit_info.total_qubits,
            physical_qubits_needed=physical_qubits,
            circuit_depth=circuit_info.depth,
            gate_count=circuit_info.gates_count,
            cnot_gates=circuit_info.cnot_count,
            single_qubit_gates=circuit_info.single_qubit_gates,
            estimated_runtime_seconds=estimated_runtime,
            estimated_error_rate=circuit_error_rate,
            success_probability=success_prob,
            qubit_quality=qubit_quality,
            error_correction_overhead=error_correction_overhead
        )


# ============================================================================
# Quantum Circuit Visualization
# ============================================================================

class QuantumCircuitVisualizer:
    """Visualize Shor's algorithm quantum circuit in ASCII art."""
    
    @staticmethod
    def draw_circuit_ascii(key_size: int = 8) -> str:
        """
        Draw simplified ASCII representation of Shor's algorithm circuit.
        """
        circuit_info = QuantumResourceEstimator.estimate_circuit_resources(key_size)
        
        lines = []
        lines.append("=" * 80)
        lines.append(f"SHOR'S ALGORITHM QUANTUM CIRCUIT - RSA-{key_size}")
        lines.append("=" * 80)
        lines.append("")
        
        # Circuit parameters
        lines.append("CIRCUIT PARAMETERS:")
        lines.append(f"  Total Qubits: {circuit_info.total_qubits}")
        lines.append(f"  ├─ Control Qubits (superposition): {circuit_info.control_qubits}")
        lines.append(f"  ├─ Target Qubits (register): {circuit_info.target_qubits}")
        lines.append(f"  └─ Ancilla Qubits (auxiliary): {circuit_info.ancilla_qubits}")
        lines.append("")
        
        # Circuit depth and gates
        lines.append("CIRCUIT COMPLEXITY:")
        lines.append(f"  Circuit Depth: {circuit_info.depth}")
        lines.append(f"  Total Gates: {circuit_info.gates_count}")
        lines.append(f"  ├─ CNOT Gates: {circuit_info.cnot_count}")
        lines.append(f"  └─ Single-qubit Gates: {circuit_info.single_qubit_gates}")
        lines.append("")
        
        # Circuit structure (ASCII art)
        lines.append("QUANTUM CIRCUIT STRUCTURE:")
        lines.append("")
        lines.append("  ┌────────────────────────────────────────────┐")
        lines.append("  │ INITIALIZATION STAGE                       │")
        lines.append("  │ Apply Hadamard gates to create superposition│")
        lines.append("  └────────────────────────────────────────────┘")
        lines.append("           |                |")
        lines.append("  ┌────────▼─────────────────▼────────────────┐")
        lines.append("  │ QUANTUM FOURIER TRANSFORM (QFT)          │")
        lines.append("  │ Order-finding using quantum phase estimation│")
        lines.append("  └────────────┬────────────────────────────────┘")
        lines.append("               |")
        lines.append("  ┌────────────▼────────────────────────────────┐")
        lines.append("  │ MEASUREMENT STAGE                          │")
        lines.append("  │ Measure output register                    │")
        lines.append("  └────────────┬────────────────────────────────┘")
        lines.append("               |")
        lines.append("  ┌────────────▼────────────────────────────────┐")
        lines.append("  │ CLASSICAL POST-PROCESSING                 │")
        lines.append("  │ Extract factors from measurement result   │")
        lines.append("  └────────────────────────────────────────────┘")
        lines.append("")
        
        return "\n".join(lines)
    
    @staticmethod
    def draw_algorithm_flow() -> str:
        """Draw the overall algorithm flow."""
        lines = []
        lines.append("=" * 80)
        lines.append("SHOR'S ALGORITHM EXECUTION FLOW")
        lines.append("=" * 80)
        lines.append("")
        
        flow = [
            ("INPUT", "RSA modulus N (to factor)"),
            ("↓", ""),
            ("CHECK", "Is N even? → Yes → Return factors (2, N/2)"),
            ("↓", ""),
            ("CHECK", "Is N a prime power? → Yes → Recursively factor"),
            ("↓", ""),
            ("LOOP", "Repeat until factored:"),
            ("  1", "Choose random 1 < a < N"),
            ("  2", "Compute gcd(a, N) → If > 1, found factor"),
            ("  3", "QUANTUM: Find order r of a modulo N"),
            ("     → Apply Hadamard gates for superposition"),
            ("     → Use Quantum Fourier Transform for period"),
            ("     → Measure output (probability > 40% for odd r)"),
            ("  4", "CLASSICAL: If r is even, extract factors"),
            ("     → Compute x = a^(r/2) mod N"),
            ("     → Compute factors as gcd(x±1, N)"),
            ("  5", "Verify: Check if p * q == N"),
            ("↓", ""),
            ("OUTPUT", "Prime factors p and q"),
        ]
        
        for label, desc in flow:
            if label in ["INPUT", "OUTPUT", "↓"]:
                lines.append(f"{label}".center(80))
            elif label == "CHECK":
                lines.append(f"  ┌─ {desc}")
            elif label == "LOOP":
                lines.append(f"  ┌─ {desc}")
            elif label.startswith("  "):
                lines.append(f"  {label} {desc}")
            else:
                lines.append(f"  {label}: {desc}")
        
        lines.append("")
        return "\n".join(lines)


# ============================================================================
# Success Probability Calculator
# ============================================================================

class SuccessProbabilityCalculator:
    """Calculate success probability for Shor's algorithm."""
    
    @staticmethod
    def calculate_probability(key_size: int, qubit_quality: QubitQuality, 
                             num_attempts: int = 1) -> float:
        """
        Calculate success probability for Shor's algorithm.
        
        Theoretical: ~40% chance per attempt with perfect quantum computer
        In practice: Depends heavily on qubit quality and error correction
        """
        # Base probability (40% for odd period with even denominator)
        base_probability = 0.40
        
        # Qubit quality factors
        quality_factors = {
            QubitQuality.PERFECT: 1.0,
            QubitQuality.EXCELLENT: 0.95,
            QubitQuality.VERY_GOOD: 0.85,
            QubitQuality.GOOD: 0.70,
            QubitQuality.FAIR: 0.50,
            QubitQuality.POOR: 0.20,
        }
        quality_factor = quality_factors.get(qubit_quality, 0.5)
        
        # Single attempt probability
        single_attempt_prob = base_probability * quality_factor
        
        # Probability of success in at least one of num_attempts
        # P(success in n attempts) = 1 - (1 - p)^n
        success_prob = 1 - (1 - single_attempt_prob) ** num_attempts
        
        return min(success_prob, 0.99)  # Cap at 99%
    
    @staticmethod
    def calculate_expected_attempts(success_probability: float) -> int:
        """Calculate expected number of attempts for given success probability."""
        if success_probability <= 0:
            return float('inf')
        if success_probability >= 0.99:
            return 1
        
        # n = ln(1 - p_target) / ln(1 - p_single)
        single_prob = 0.40  # Base 40% probability
        expected = math.log(1 - success_probability) / math.log(1 - single_prob)
        return max(1, int(math.ceil(expected)))
    
    @staticmethod
    def get_probability_table(key_size: int) -> Dict[str, float]:
        """Get success probability for all qubit quality levels."""
        probabilities = {}
        for quality in QubitQuality:
            prob = SuccessProbabilityCalculator.calculate_probability(
                key_size, quality, num_attempts=5
            )
            probabilities[quality.value] = prob
        return probabilities


# ============================================================================
# Integration Function
# ============================================================================

def create_shors_algorithm(algorithm_id: str = None) -> ShorsAlgorithm:
    """Factory function to create Shor's algorithm instance."""
    return ShorsAlgorithm(algorithm_id)
