import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pandas as pd
from typing import List, Tuple, Dict

class MPCProtocol:
    """Base class for MPC protocols"""
    def __init__(self, security_param: int):
        self.security_param = security_param
    
    def comm_cost_per_gate(self, num_parties: int) -> float:
        """Communication cost per multiplication gate"""
        raise NotImplementedError

class BeaverProtocol(MPCProtocol):
    """Beaver triple-based protocol (GMW-style)"""
    def comm_cost_per_gate(self, num_parties: int) -> float:
        # Each party broadcasts one value per gate
        # Cost: O(n^2 * lambda) bits
        return num_parties * (num_parties - 1) * self.security_param / 8  # Convert to bytes

class DNNProtocol(MPCProtocol):
    """Dealer-based protocol with preprocessing"""
    def comm_cost_per_gate(self, num_parties: int) -> float:
        # Dealer sends to each party, parties send back
        # Cost: O(n * lambda) bits
        return 2 * num_parties * self.security_param / 8

class SPDZProtocol(MPCProtocol):
    """SPDZ-style protocol with MAC"""
    def comm_cost_per_gate(self, num_parties: int) -> float:
        # Each party broadcasts value + MAC
        # Cost: O(n^2 * 2*lambda) bits
        return num_parties * (num_parties - 1) * 2 * self.security_param / 8

class HybridProtocol(MPCProtocol):
    """Optimized hybrid protocol"""
    def comm_cost_per_gate(self, num_parties: int) -> float:
        # Uses threshold optimization
        if num_parties <= 3:
            return num_parties * self.security_param / 8
        else:
            # Hierarchical approach for larger groups
            return (2 * np.sqrt(num_parties) * self.security_param / 8 + 
                    num_parties * 0.5 * self.security_param / 8)

def compute_circuit_complexity(num_values: int) -> int:
    """
    Calculate number of multiplication gates for sum of squared differences
    Formula: sum((x_i - mean)^2) requires multiplications
    """
    # Squaring each difference: num_values multiplications
    # Computing mean involves additions only (no multiplications in MPC sense)
    return num_values

def optimize_communication_cost(
    num_parties_range: List[int],
    num_values_range: List[int],
    security_param: int = 128
) -> Dict:
    """
    Optimize communication cost across different protocols
    """
    protocols = {
        'Beaver': BeaverProtocol(security_param),
        'DNN': DNNProtocol(security_param),
        'SPDZ': SPDZProtocol(security_param),
        'Hybrid': HybridProtocol(security_param)
    }
    
    results = {name: np.zeros((len(num_parties_range), len(num_values_range))) 
               for name in protocols.keys()}
    
    for i, num_parties in enumerate(num_parties_range):
        for j, num_values in enumerate(num_values_range):
            num_gates = compute_circuit_complexity(num_values)
            
            for name, protocol in protocols.items():
                cost_per_gate = protocol.comm_cost_per_gate(num_parties)
                total_cost = cost_per_gate * num_gates
                results[name][i, j] = total_cost / 1024  # Convert to KB
    
    return results, protocols

def find_optimal_protocol(
    num_parties: int,
    num_values: int,
    protocols: Dict
) -> Tuple[str, float]:
    """Find the protocol with minimum communication cost"""
    num_gates = compute_circuit_complexity(num_values)
    
    min_cost = float('inf')
    best_protocol = None
    
    for name, protocol in protocols.items():
        cost = protocol.comm_cost_per_gate(num_parties) * num_gates / 1024
        if cost < min_cost:
            min_cost = cost
            best_protocol = name
    
    return best_protocol, min_cost

def analyze_optimization_ratio(results: Dict, num_parties_range: List[int]) -> pd.DataFrame:
    """Analyze optimization ratios compared to baseline"""
    baseline = results['Beaver']
    
    ratios = {}
    for name, cost_matrix in results.items():
        if name != 'Beaver':
            ratio = (baseline - cost_matrix) / baseline * 100
            ratios[name] = ratio.mean()
    
    df = pd.DataFrame({
        'Protocol': list(ratios.keys()),
        'Avg Reduction (%)': list(ratios.values())
    })
    return df.sort_values('Avg Reduction (%)', ascending=False)

# Main execution
print("=" * 70)
print("MPC Communication Cost Optimization Analysis")
print("=" * 70)
print()

# Parameters
security_param = 128
num_parties_range = list(range(2, 21))
num_values_range = list(range(10, 201, 10))

print(f"Security Parameter (λ): {security_param} bits")
print(f"Number of Parties Range: {num_parties_range[0]} to {num_parties_range[-1]}")
print(f"Number of Values Range: {num_values_range[0]} to {num_values_range[-1]}")
print()

# Compute optimization results
print("Computing communication costs for different protocols...")
results, protocols = optimize_communication_cost(
    num_parties_range,
    num_values_range,
    security_param
)
print("Done!")
print()

# Analyze specific cases
print("=" * 70)
print("Case Study: Optimal Protocol Selection")
print("=" * 70)
test_cases = [(3, 50), (5, 100), (10, 150), (15, 200)]

for num_parties, num_values in test_cases:
    best_protocol, min_cost = find_optimal_protocol(num_parties, num_values, protocols)
    print(f"Parties: {num_parties:2d}, Values: {num_values:3d} -> "
          f"Best: {best_protocol:8s} ({min_cost:8.2f} KB)")
print()

# Optimization ratios
print("=" * 70)
print("Communication Reduction vs Baseline (Beaver Protocol)")
print("=" * 70)
ratio_df = analyze_optimization_ratio(results, num_parties_range)
print(ratio_df.to_string(index=False))
print()

# Visualization
fig = plt.figure(figsize=(18, 12))

# Plot 1: 3D Surface for Hybrid Protocol
ax1 = fig.add_subplot(2, 3, 1, projection='3d')
X, Y = np.meshgrid(num_values_range, num_parties_range)
surf1 = ax1.plot_surface(X, Y, results['Hybrid'], cmap=cm.viridis, alpha=0.8)
ax1.set_xlabel('Number of Values', fontsize=10)
ax1.set_ylabel('Number of Parties', fontsize=10)
ax1.set_zlabel('Communication Cost (KB)', fontsize=10)
ax1.set_title('Hybrid Protocol: 3D Communication Cost', fontsize=11, fontweight='bold')
fig.colorbar(surf1, ax=ax1, shrink=0.5)

# Plot 2: Comparison of all protocols (2D heatmap style)
ax2 = fig.add_subplot(2, 3, 2)
for name, cost_matrix in results.items():
    # Average over number of values for each party count
    avg_costs = cost_matrix.mean(axis=1)
    ax2.plot(num_parties_range, avg_costs, marker='o', label=name, linewidth=2)
ax2.set_xlabel('Number of Parties', fontsize=11)
ax2.set_ylabel('Avg Communication Cost (KB)', fontsize=11)
ax2.set_title('Protocol Comparison: Communication vs Parties', fontsize=11, fontweight='bold')
ax2.legend()
ax2.grid(True, alpha=0.3)

# Plot 3: Cost scaling with values
ax3 = fig.add_subplot(2, 3, 3)
fixed_parties = [3, 7, 12, 18]
for np_val in fixed_parties:
    idx = num_parties_range.index(np_val)
    ax3.plot(num_values_range, results['Hybrid'][idx, :], 
             marker='s', label=f'{np_val} parties', linewidth=2)
ax3.set_xlabel('Number of Values', fontsize=11)
ax3.set_ylabel('Communication Cost (KB)', fontsize=11)
ax3.set_title('Hybrid Protocol: Scaling with Circuit Size', fontsize=11, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)

# Plot 4: 3D comparison (Beaver vs Hybrid)
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
reduction = (results['Beaver'] - results['Hybrid']) / results['Beaver'] * 100
surf2 = ax4.plot_surface(X, Y, reduction, cmap=cm.RdYlGn, alpha=0.8)
ax4.set_xlabel('Number of Values', fontsize=10)
ax4.set_ylabel('Number of Parties', fontsize=10)
ax4.set_zlabel('Reduction (%)', fontsize=10)
ax4.set_title('Communication Reduction: Hybrid vs Beaver', fontsize=11, fontweight='bold')
fig.colorbar(surf2, ax=ax4, shrink=0.5)

# Plot 5: Protocol efficiency ratio
ax5 = fig.add_subplot(2, 3, 5)
protocol_names = list(results.keys())
total_costs = [results[name].sum() for name in protocol_names]
colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4']
bars = ax5.bar(protocol_names, total_costs, color=colors, alpha=0.8, edgecolor='black')
ax5.set_ylabel('Total Communication Cost (KB)', fontsize=11)
ax5.set_title('Total Communication Cost by Protocol', fontsize=11, fontweight='bold')
ax5.grid(True, alpha=0.3, axis='y')
for bar in bars:
    height = bar.get_height()
    ax5.text(bar.get_x() + bar.get_width()/2., height,
             f'{int(height)}', ha='center', va='bottom', fontsize=9)

# Plot 6: Optimal protocol selection heatmap
ax6 = fig.add_subplot(2, 3, 6)
optimal_matrix = np.zeros((len(num_parties_range), len(num_values_range)))
for i, num_parties in enumerate(num_parties_range):
    for j, num_values in enumerate(num_values_range):
        min_cost = float('inf')
        best_idx = 0
        for idx, name in enumerate(results.keys()):
            if results[name][i, j] < min_cost:
                min_cost = results[name][i, j]
                best_idx = idx
        optimal_matrix[i, j] = best_idx

im = ax6.imshow(optimal_matrix, cmap='Set3', aspect='auto', origin='lower')
ax6.set_xlabel('Number of Values (index)', fontsize=11)
ax6.set_ylabel('Number of Parties (index)', fontsize=11)
ax6.set_title('Optimal Protocol Selection Map', fontsize=11, fontweight='bold')
cbar = fig.colorbar(im, ax=ax6, ticks=[0, 1, 2, 3])
cbar.set_ticklabels(list(results.keys()))

plt.tight_layout()
plt.savefig('mpc_communication_optimization.png', dpi=300, bbox_inches='tight')
print("Visualization saved as 'mpc_communication_optimization.png'")
print()

# Additional 3D visualization
fig2 = plt.figure(figsize=(16, 6))

# 3D plot for all protocols
for idx, (name, cost_matrix) in enumerate(results.items()):
    ax = fig2.add_subplot(1, 4, idx + 1, projection='3d')
    surf = ax.plot_surface(X, Y, cost_matrix, cmap=cm.coolwarm, alpha=0.8)
    ax.set_xlabel('Values', fontsize=9)
    ax.set_ylabel('Parties', fontsize=9)
    ax.set_zlabel('Cost (KB)', fontsize=9)
    ax.set_title(f'{name} Protocol', fontsize=10, fontweight='bold')
    ax.view_init(elev=25, azim=45)
    fig2.colorbar(surf, ax=ax, shrink=0.6)

plt.tight_layout()
plt.savefig('mpc_3d_comparison.png', dpi=300, bbox_inches='tight')
print("3D comparison saved as 'mpc_3d_comparison.png'")
print()

print("=" * 70)
print("Analysis Complete!")
print("=" * 70)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import hashlib
import time
import secrets
from dataclasses import dataclass
from typing import List, Tuple

# Cryptographic parameters
def generate_safe_prime_params():
    """Generate safe cryptographic parameters for demonstration"""
    # Using a 128-bit safe prime for practical demonstration
    # In production, use 2048+ bit primes
    p = 2**127 - 1  # Mersenne prime
    g = 2
    q = (p - 1) // 2
    return p, g, q

def mod_exp(base, exp, mod):
    """Fast modular exponentiation"""
    return pow(int(base), int(exp), int(mod))

def random_in_range(upper_bound):
    """Generate cryptographically secure random number in range [1, upper_bound)"""
    # Use secrets module for cryptographic randomness
    # Handle large numbers by using bytes
    if upper_bound <= 2**63:
        return secrets.randbelow(upper_bound - 1) + 1
    else:
        # For large numbers, use byte generation
        num_bytes = (upper_bound.bit_length() + 7) // 8
        while True:
            random_bytes = secrets.token_bytes(num_bytes)
            random_num = int.from_bytes(random_bytes, 'big')
            if 1 <= random_num < upper_bound:
                return random_num

@dataclass
class ProofMetrics:
    """Metrics for proof communication complexity"""
    rounds: int
    proof_size_bits: int
    computation_time: float
    soundness_error: float

class InteractiveSchnorrProtocol:
    """Interactive Schnorr protocol for discrete log"""
    
    def __init__(self, p: int, g: int, q: int, security_param: int = 128):
        self.p = p
        self.g = g
        self.q = q
        self.security_param = security_param
        self.repetitions = security_param  # Number of parallel repetitions
        
    def prove_interactive(self, x: int, h: int) -> ProofMetrics:
        """Interactive proof with multiple rounds"""
        start_time = time.time()
        
        # Single round requires log2(q) bits for challenge
        # To achieve 2^-lambda security, we need lambda repetitions
        total_commitment_bits = self.repetitions * 128  # Each commitment is ~128 bits
        total_challenge_bits = self.repetitions * 1  # Each challenge is 1 bit for basic
        total_response_bits = self.repetitions * 128  # Each response is ~128 bits
        
        # Simulate the protocol
        for _ in range(self.repetitions):
            # Prover: commit with random r
            r = random_in_range(self.q)
            commitment = mod_exp(self.g, r, self.p)
            
            # Verifier: send challenge (binary for simplicity)
            challenge = secrets.randbelow(2)
            
            # Prover: compute response
            response = (r + challenge * x) % self.q
            
            # Verifier: check
            left = mod_exp(self.g, response, self.p)
            right = (commitment * mod_exp(h, challenge, self.p)) % self.p
            assert left == right, "Verification failed"
        
        computation_time = time.time() - start_time
        
        # Each repetition requires 3 messages (commitment, challenge, response)
        # But we can parallelize: send all commitments, receive all challenges, send all responses
        rounds = 3
        proof_size = total_commitment_bits + total_challenge_bits + total_response_bits
        soundness_error = 2 ** (-self.repetitions)
        
        return ProofMetrics(rounds, proof_size, computation_time, soundness_error)
    
    def prove_non_interactive(self, x: int, h: int) -> ProofMetrics:
        """Non-interactive proof using Fiat-Shamir heuristic"""
        start_time = time.time()
        
        # Generate random nonce
        r = random_in_range(self.q)
        commitment = mod_exp(self.g, r, self.p)
        
        # Fiat-Shamir: challenge = Hash(g, h, commitment)
        hash_input = f"{self.g}{h}{commitment}".encode()
        challenge_hash = hashlib.sha256(hash_input).digest()
        challenge = int.from_bytes(challenge_hash, 'big') % self.q
        
        # Compute response
        response = (r + challenge * x) % self.q
        
        # Verify
        left = mod_exp(self.g, response, self.p)
        right = (commitment * mod_exp(h, challenge, self.p)) % self.p
        assert left == right, "Verification failed"
        
        computation_time = time.time() - start_time
        
        # Non-interactive: only 1 round (prover sends proof)
        rounds = 1
        # Proof = (commitment, response) - challenge is derived
        proof_size = 128 + 128  # commitment + response in bits
        soundness_error = 2 ** (-128)  # Security of hash function
        
        return ProofMetrics(rounds, proof_size, computation_time, soundness_error)

class OptimizedBatchProtocol:
    """Optimized protocol with batching and compression"""
    
    def __init__(self, p: int, g: int, q: int, security_param: int = 128):
        self.p = p
        self.g = g
        self.q = q
        self.security_param = security_param
        
    def prove_batch(self, secrets: List[int], publics: List[int]) -> ProofMetrics:
        """Batch proof for multiple statements"""
        start_time = time.time()
        n = len(secrets)
        
        # Single commitment for all statements using random linear combination
        r = random_in_range(self.q)
        commitment = mod_exp(self.g, r, self.p)
        
        # Fiat-Shamir challenge
        hash_input = f"{self.g}{''.join(map(str, publics))}{commitment}".encode()
        challenge = int.from_bytes(hashlib.sha256(hash_input).digest(), 'big') % self.q
        
        # Aggregated response
        response = r
        for i, x in enumerate(secrets):
            weight = mod_exp(challenge, i + 1, self.q)
            response = (response + weight * x) % self.q
        
        computation_time = time.time() - start_time
        
        # Single commitment and single response regardless of n
        rounds = 1
        proof_size = 128 + 128  # Constant size!
        soundness_error = 2 ** (-128)
        
        return ProofMetrics(rounds, proof_size, computation_time, soundness_error)

def compare_protocols():
    """Compare different protocol variants"""
    # Setup
    p, g, q = generate_safe_prime_params()
    
    # Generate secret and public values
    x = random_in_range(q)
    h = mod_exp(g, x, p)
    
    # Test different security parameters
    security_params = [32, 64, 96, 128, 160, 192, 224, 256]
    
    interactive_results = []
    non_interactive_results = []
    
    print("Testing Interactive Protocol...")
    for sec_param in security_params:
        protocol = InteractiveSchnorrProtocol(p, g, q, sec_param)
        metrics = protocol.prove_interactive(x, h)
        interactive_results.append(metrics)
        print(f"Security={sec_param}: Rounds={metrics.rounds}, Size={metrics.proof_size_bits} bits, Time={metrics.computation_time:.4f}s")
    
    print("\nTesting Non-Interactive Protocol...")
    for sec_param in security_params:
        protocol = InteractiveSchnorrProtocol(p, g, q, sec_param)
        metrics = protocol.prove_non_interactive(x, h)
        non_interactive_results.append(metrics)
        print(f"Security={sec_param}: Rounds={metrics.rounds}, Size={metrics.proof_size_bits} bits, Time={metrics.computation_time:.4f}s")
    
    # Test batch protocol
    print("\nTesting Batch Protocol...")
    batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
    batch_results = []
    
    for batch_size in batch_sizes:
        secrets = [random_in_range(q) for _ in range(batch_size)]
        publics = [mod_exp(g, secret, p) for secret in secrets]
        
        protocol = OptimizedBatchProtocol(p, g, q, 128)
        metrics = protocol.prove_batch(secrets, publics)
        batch_results.append(metrics)
        print(f"Batch size={batch_size}: Size={metrics.proof_size_bits} bits, Time={metrics.computation_time:.4f}s")
    
    return security_params, interactive_results, non_interactive_results, batch_sizes, batch_results

def visualize_results(security_params, interactive_results, non_interactive_results, 
                     batch_sizes, batch_results):
    """Create comprehensive visualizations"""
    
    fig = plt.figure(figsize=(20, 12))
    
    # Plot 1: Proof Size vs Security Parameter
    ax1 = fig.add_subplot(2, 3, 1)
    interactive_sizes = [m.proof_size_bits for m in interactive_results]
    non_interactive_sizes = [m.proof_size_bits for m in non_interactive_results]
    
    ax1.plot(security_params, interactive_sizes, 'o-', label='Interactive', linewidth=2, markersize=8)
    ax1.plot(security_params, non_interactive_sizes, 's-', label='Non-Interactive (Fiat-Shamir)', linewidth=2, markersize=8)
    ax1.set_xlabel('Security Parameter λ (bits)', fontsize=12)
    ax1.set_ylabel('Proof Size (bits)', fontsize=12)
    ax1.set_title('Communication Complexity vs Security', fontsize=14, fontweight='bold')
    ax1.legend(fontsize=10)
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Rounds Comparison
    ax2 = fig.add_subplot(2, 3, 2)
    interactive_rounds = [m.rounds for m in interactive_results]
    non_interactive_rounds = [m.rounds for m in non_interactive_results]
    
    ax2.bar(np.arange(len(security_params)) - 0.2, interactive_rounds, 0.4, label='Interactive', alpha=0.8)
    ax2.bar(np.arange(len(security_params)) + 0.2, non_interactive_rounds, 0.4, label='Non-Interactive', alpha=0.8)
    ax2.set_xlabel('Security Parameter λ (bits)', fontsize=12)
    ax2.set_ylabel('Number of Rounds', fontsize=12)
    ax2.set_title('Round Complexity Comparison', fontsize=14, fontweight='bold')
    ax2.set_xticks(np.arange(len(security_params)))
    ax2.set_xticklabels(security_params, rotation=45)
    ax2.legend(fontsize=10)
    ax2.grid(True, alpha=0.3, axis='y')
    
    # Plot 3: Computation Time
    ax3 = fig.add_subplot(2, 3, 3)
    interactive_times = [m.computation_time * 1000 for m in interactive_results]
    non_interactive_times = [m.computation_time * 1000 for m in non_interactive_results]
    
    ax3.semilogy(security_params, interactive_times, 'o-', label='Interactive', linewidth=2, markersize=8)
    ax3.semilogy(security_params, non_interactive_times, 's-', label='Non-Interactive', linewidth=2, markersize=8)
    ax3.set_xlabel('Security Parameter λ (bits)', fontsize=12)
    ax3.set_ylabel('Computation Time (ms)', fontsize=12)
    ax3.set_title('Computational Overhead', fontsize=14, fontweight='bold')
    ax3.legend(fontsize=10)
    ax3.grid(True, alpha=0.3)
    
    # Plot 4: Batch Efficiency
    ax4 = fig.add_subplot(2, 3, 4)
    batch_sizes_list = batch_sizes
    batch_proof_sizes = [m.proof_size_bits for m in batch_results]
    naive_sizes = [256 * size for size in batch_sizes_list]  # 256 bits per proof naively
    
    ax4.plot(batch_sizes_list, naive_sizes, 'o-', label='Naive (No Batching)', linewidth=2, markersize=8)
    ax4.plot(batch_sizes_list, batch_proof_sizes, 's-', label='Optimized Batch', linewidth=2, markersize=8)
    ax4.set_xlabel('Number of Statements', fontsize=12)
    ax4.set_ylabel('Total Proof Size (bits)', fontsize=12)
    ax4.set_title('Batch Proof Efficiency', fontsize=14, fontweight='bold')
    ax4.legend(fontsize=10)
    ax4.grid(True, alpha=0.3)
    ax4.set_xscale('log', base=2)
    
    # Plot 5: Soundness Error
    ax5 = fig.add_subplot(2, 3, 5)
    interactive_errors = [-np.log2(m.soundness_error) for m in interactive_results]
    non_interactive_errors = [-np.log2(m.soundness_error) for m in non_interactive_results]
    
    ax5.plot(security_params, interactive_errors, 'o-', label='Interactive', linewidth=2, markersize=8)
    ax5.plot(security_params, non_interactive_errors, 's-', label='Non-Interactive', linewidth=2, markersize=8)
    ax5.plot(security_params, security_params, '--', label='Target Security', linewidth=2, alpha=0.5)
    ax5.set_xlabel('Security Parameter λ (bits)', fontsize=12)
    ax5.set_ylabel('Achieved Security -log₂(ε)', fontsize=12)
    ax5.set_title('Soundness Error Analysis', fontsize=14, fontweight='bold')
    ax5.legend(fontsize=10)
    ax5.grid(True, alpha=0.3)
    
    # Plot 6: 3D Trade-off Surface
    ax6 = fig.add_subplot(2, 3, 6, projection='3d')
    
    # Create mesh for 3D plot
    X, Y = np.meshgrid(security_params, [1, 3])  # Rounds dimension
    Z_size = np.zeros_like(X, dtype=float)
    
    for i, sec_param in enumerate(security_params):
        Z_size[0, i] = non_interactive_results[i].proof_size_bits
        Z_size[1, i] = interactive_results[i].proof_size_bits
    
    surf = ax6.plot_surface(X, Y, Z_size, cmap='viridis', alpha=0.8, edgecolor='none')
    ax6.set_xlabel('Security λ (bits)', fontsize=10)
    ax6.set_ylabel('Rounds', fontsize=10)
    ax6.set_zlabel('Proof Size (bits)', fontsize=10)
    ax6.set_title('3D Trade-off: Security vs Rounds vs Size', fontsize=12, fontweight='bold')
    fig.colorbar(surf, ax=ax6, shrink=0.5)
    
    plt.tight_layout()
    plt.savefig('zkp_communication_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # Additional 3D visualization: Batch performance
    fig2 = plt.figure(figsize=(14, 6))
    
    ax7 = fig2.add_subplot(1, 2, 1, projection='3d')
    batch_x = np.array(batch_sizes_list)
    batch_y = np.array([m.proof_size_bits for m in batch_results])
    batch_z = np.array([m.computation_time * 1000 for m in batch_results])
    
    ax7.scatter(batch_x, batch_y, batch_z, c=batch_z, cmap='plasma', s=100, alpha=0.8)
    ax7.plot(batch_x, batch_y, batch_z, 'b-', alpha=0.5, linewidth=2)
    ax7.set_xlabel('Batch Size', fontsize=11)
    ax7.set_ylabel('Proof Size (bits)', fontsize=11)
    ax7.set_zlabel('Time (ms)', fontsize=11)
    ax7.set_title('Batch Protocol: Size-Time Trade-off', fontsize=13, fontweight='bold')
    
    ax8 = fig2.add_subplot(1, 2, 2, projection='3d')
    
    # Create efficiency surface
    batch_array = np.array(batch_sizes_list)
    efficiency_naive = batch_array * 256  # bits per statement
    efficiency_batch = np.array([m.proof_size_bits for m in batch_results])
    compression_ratio = efficiency_naive / efficiency_batch
    
    X_batch = np.tile(batch_array.reshape(-1, 1), (1, 2))
    Y_methods = np.tile(np.array([0, 1]).reshape(1, -1), (len(batch_array), 1))
    Z_efficiency = np.column_stack([efficiency_batch, efficiency_naive])
    
    for i in range(len(batch_array)):
        ax8.plot([batch_array[i], batch_array[i]], [0, 1], 
                [efficiency_batch[i], efficiency_naive[i]], 'gray', alpha=0.3)
    
    ax8.scatter(X_batch[:, 0], Y_methods[:, 0], Z_efficiency[:, 0], 
               c='green', s=100, label='Batch', alpha=0.8)
    ax8.scatter(X_batch[:, 1], Y_methods[:, 1], Z_efficiency[:, 1], 
               c='red', s=100, label='Naive', alpha=0.8)
    
    ax8.set_xlabel('Number of Statements', fontsize=11)
    ax8.set_ylabel('Method', fontsize=11)
    ax8.set_zlabel('Total Proof Size (bits)', fontsize=11)
    ax8.set_yticks([0, 1])
    ax8.set_yticklabels(['Batch', 'Naive'])
    ax8.set_title('Batching Compression Advantage', fontsize=13, fontweight='bold')
    ax8.legend(fontsize=10)
    
    plt.tight_layout()
    plt.savefig('zkp_batch_analysis_3d.png', dpi=300, bbox_inches='tight')
    plt.show()

def print_summary_table(security_params, interactive_results, non_interactive_results):
    """Print detailed comparison table"""
    print("\n" + "="*100)
    print("COMPREHENSIVE COMPARISON TABLE")
    print("="*100)
    print(f"{'Security λ':<12} {'Protocol':<20} {'Rounds':<8} {'Size (bits)':<12} {'Time (ms)':<12} {'Soundness':<15}")
    print("-"*100)
    
    for i, sec_param in enumerate(security_params):
        inter = interactive_results[i]
        non_inter = non_interactive_results[i]
        
        print(f"{sec_param:<12} {'Interactive':<20} {inter.rounds:<8} {inter.proof_size_bits:<12} "
              f"{inter.computation_time*1000:<12.4f} {inter.soundness_error:<15.2e}")
        print(f"{'':<12} {'Non-Interactive':<20} {non_inter.rounds:<8} {non_inter.proof_size_bits:<12} "
              f"{non_inter.computation_time*1000:<12.4f} {non_inter.soundness_error:<15.2e}")
        print("-"*100)

# Main execution
print("Zero-Knowledge Proof Communication Complexity Analysis")
print("="*60)
print("\nInitializing cryptographic parameters...")

security_params, interactive_results, non_interactive_results, batch_sizes, batch_results = compare_protocols()

print_summary_table(security_params, interactive_results, non_interactive_results)

print("\n" + "="*100)
print("BATCH PROTOCOL ANALYSIS")
print("="*100)
print(f"{'Batch Size':<12} {'Proof Size (bits)':<20} {'Naive Size (bits)':<20} {'Compression Ratio':<20} {'Time (ms)':<12}")
print("-"*100)

for i, batch_size in enumerate(batch_sizes):
    naive_size = batch_size * 256
    actual_size = batch_results[i].proof_size_bits
    ratio = naive_size / actual_size
    time_ms = batch_results[i].computation_time * 1000
    print(f"{batch_size:<12} {actual_size:<20} {naive_size:<20} {ratio:<20.2f}x {time_ms:<12.4f}")

print("\n" + "="*100)
print("\nGenerating visualizations...")
visualize_results(security_params, interactive_results, non_interactive_results, batch_sizes, batch_results)

print("\nAnalysis complete! Graphs saved as 'zkp_communication_analysis.png' and 'zkp_batch_analysis_3d.png'")

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from pulp import *
import itertools

# Define a simplified 4-bit S-box with differential distribution table (DDT)
SBOX = [0xC, 0x5, 0x6, 0xB, 0x9, 0x0, 0xA, 0xD, 0x3, 0xE, 0xF, 0x8, 0x4, 0x7, 0x1, 0x2]

def compute_ddt(sbox):
    """Compute Differential Distribution Table for given S-box"""
    n = len(sbox)
    ddt = np.zeros((n, n), dtype=int)
    
    for input_diff in range(n):
        for x in range(n):
            output_diff = sbox[x] ^ sbox[x ^ input_diff]
            ddt[input_diff][output_diff] += 1
    
    return ddt

def ddt_to_probability(ddt):
    """Convert DDT counts to probabilities"""
    n = ddt.shape[0]
    prob_table = ddt.astype(float) / n
    return prob_table

def compute_log_probability(prob):
    """Compute -log2(probability), handling zero probabilities"""
    with np.errstate(divide='ignore'):
        log_prob = -np.log2(prob)
        log_prob[np.isinf(log_prob)] = 1000  # Large penalty for impossible differentials
    return log_prob

# Build DDT and probability tables
DDT = compute_ddt(SBOX)
PROB_TABLE = ddt_to_probability(DDT)
LOG_PROB_TABLE = compute_log_probability(PROB_TABLE)

print("Differential Distribution Table (DDT):")
print(DDT)
print("\nProbability Table (first 8x8):")
print(PROB_TABLE[:8, :8])

# MILP Model for finding optimal differential path
def find_optimal_differential_path(num_rounds=4, num_sboxes=4):
    """
    Find optimal differential path using MILP
    
    Parameters:
    - num_rounds: Number of cipher rounds
    - num_sboxes: Number of parallel S-boxes per round
    """
    
    # Create MILP problem
    prob = LpProblem("Optimal_Differential_Path", LpMinimize)
    
    # Decision variables: binary variables for each possible differential at each S-box
    # delta[round][sbox][input_diff][output_diff]
    delta = {}
    for r in range(num_rounds):
        for s in range(num_sboxes):
            for i_diff in range(16):
                for o_diff in range(16):
                    var_name = f"d_r{r}_s{s}_i{i_diff}_o{o_diff}"
                    delta[(r, s, i_diff, o_diff)] = LpVariable(var_name, cat='Binary')
    
    # Objective: minimize total -log2(probability)
    objective = []
    for r in range(num_rounds):
        for s in range(num_sboxes):
            for i_diff in range(16):
                for o_diff in range(16):
                    cost = LOG_PROB_TABLE[i_diff, o_diff]
                    if cost < 1000:  # Only add feasible transitions
                        objective.append(cost * delta[(r, s, i_diff, o_diff)])
    
    prob += lpSum(objective)
    
    # Constraints
    # 1. Each S-box in each round must have exactly one active differential
    for r in range(num_rounds):
        for s in range(num_sboxes):
            prob += lpSum([delta[(r, s, i_diff, o_diff)] 
                          for i_diff in range(16) 
                          for o_diff in range(16)]) == 1
    
    # 2. Eliminate impossible differentials (probability = 0)
    for r in range(num_rounds):
        for s in range(num_sboxes):
            for i_diff in range(16):
                for o_diff in range(16):
                    if DDT[i_diff, o_diff] == 0:
                        prob += delta[(r, s, i_diff, o_diff)] == 0
    
    # 3. Simple linear layer: output of round r connects to input of round r+1
    # For simplicity, we use bit permutation: sbox i connects to sbox (i+1) mod num_sboxes
    for r in range(num_rounds - 1):
        for s in range(num_sboxes):
            next_s = (s + 1) % num_sboxes
            for diff in range(16):
                # If output diff of (r,s) is 'diff', then input diff of (r+1,next_s) must be 'diff'
                prob += lpSum([delta[(r, s, i_d, diff)] for i_d in range(16)]) == \
                        lpSum([delta[(r+1, next_s, diff, o_d)] for o_d in range(16)])
    
    # 4. Force at least one non-zero input differential in first round
    prob += lpSum([delta[(0, s, 0, o_diff)] 
                   for s in range(num_sboxes) 
                   for o_diff in range(16)]) <= num_sboxes - 1
    
    # Solve
    prob.solve(PULP_CBC_CMD(msg=0))
    
    # Extract solution
    path = []
    total_log_prob = 0
    
    for r in range(num_rounds):
        round_diffs = []
        for s in range(num_sboxes):
            for i_diff in range(16):
                for o_diff in range(16):
                    if value(delta[(r, s, i_diff, o_diff)]) > 0.5:
                        round_diffs.append((i_diff, o_diff))
                        if DDT[i_diff, o_diff] > 0:
                            total_log_prob += LOG_PROB_TABLE[i_diff, o_diff]
        path.append(round_diffs)
    
    probability = 2 ** (-total_log_prob)
    
    return path, probability, total_log_prob

# Solve the problem
print("\n" + "="*60)
print("Finding Optimal Differential Path...")
print("="*60)

differential_path, max_probability, total_cost = find_optimal_differential_path(num_rounds=4, num_sboxes=4)

print(f"\nOptimal Path Found!")
print(f"Total -log2(P): {total_cost:.4f}")
print(f"Path Probability: 2^({-total_cost:.4f}) = {max_probability:.6e}")

print("\nDifferential Path (Input Diff -> Output Diff for each S-box):")
for r, round_diffs in enumerate(differential_path):
    print(f"\nRound {r}:")
    for s, (in_diff, out_diff) in enumerate(round_diffs):
        prob = PROB_TABLE[in_diff, out_diff]
        print(f"  S-box {s}: 0x{in_diff:X} -> 0x{out_diff:X} (p = {prob:.4f})")

# Visualization 1: Probability distribution heatmap
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

im1 = axes[0].imshow(PROB_TABLE, cmap='YlOrRd', aspect='auto')
axes[0].set_title('S-box Differential Probability Table', fontsize=14, fontweight='bold')
axes[0].set_xlabel('Output Difference')
axes[0].set_ylabel('Input Difference')
plt.colorbar(im1, ax=axes[0], label='Probability')

# Mark impossible differentials
im2 = axes[1].imshow(DDT, cmap='Blues', aspect='auto')
axes[1].set_title('Differential Distribution Table (DDT)', fontsize=14, fontweight='bold')
axes[1].set_xlabel('Output Difference')
axes[1].set_ylabel('Input Difference')
plt.colorbar(im2, ax=axes[1], label='Count')

plt.tight_layout()
plt.savefig('ddt_analysis.png', dpi=150, bbox_inches='tight')
plt.show()

# Visualization 2: Differential path flow
fig, ax = plt.subplots(figsize=(12, 8))

num_rounds = len(differential_path)
num_sboxes = len(differential_path[0])

# Draw path
for r in range(num_rounds):
    for s in range(num_sboxes):
        in_diff, out_diff = differential_path[r][s]
        
        # Draw S-box
        x_pos = r * 3
        y_pos = s * 2
        
        prob = PROB_TABLE[in_diff, out_diff]
        color_intensity = prob
        
        rect = plt.Rectangle((x_pos, y_pos), 1, 0.8, 
                            facecolor=plt.cm.RdYlGn(color_intensity),
                            edgecolor='black', linewidth=2)
        ax.add_patch(rect)
        
        # Add labels
        ax.text(x_pos + 0.5, y_pos + 0.4, f'{in_diff:X}→{out_diff:X}\n{prob:.3f}',
                ha='center', va='center', fontsize=9, fontweight='bold')
        
        # Draw connections to next round
        if r < num_rounds - 1:
            next_s = (s + 1) % num_sboxes
            ax.arrow(x_pos + 1, y_pos + 0.4, 1.5, (next_s - s) * 2,
                    head_width=0.15, head_length=0.2, fc='blue', ec='blue', alpha=0.6)

ax.set_xlim(-0.5, num_rounds * 3)
ax.set_ylim(-0.5, num_sboxes * 2)
ax.set_xlabel('Round', fontsize=12)
ax.set_ylabel('S-box Position', fontsize=12)
ax.set_title('Optimal Differential Path Flow\n(Color: Green=High Prob, Red=Low Prob)', 
             fontsize=14, fontweight='bold')
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')

plt.tight_layout()
plt.savefig('differential_path.png', dpi=150, bbox_inches='tight')
plt.show()

# Visualization 3: 3D probability landscape
fig = plt.figure(figsize=(14, 10))
ax1 = fig.add_subplot(121, projection='3d')

X, Y = np.meshgrid(range(16), range(16))
Z = PROB_TABLE

surf = ax1.plot_surface(X, Y, Z, cmap='viridis', alpha=0.8, edgecolor='none')
ax1.set_xlabel('Output Difference', fontsize=10)
ax1.set_ylabel('Input Difference', fontsize=10)
ax1.set_zlabel('Probability', fontsize=10)
ax1.set_title('3D Differential Probability Landscape', fontsize=12, fontweight='bold')
ax1.view_init(elev=25, azim=45)
fig.colorbar(surf, ax=ax1, shrink=0.5, aspect=5)

# Mark optimal path differentials
for r, round_diffs in enumerate(differential_path):
    for in_diff, out_diff in round_diffs:
        prob = PROB_TABLE[in_diff, out_diff]
        ax1.scatter([out_diff], [in_diff], [prob], c='red', s=100, marker='o', 
                   edgecolors='black', linewidth=2, zorder=10)

# 3D bar chart of round probabilities
ax2 = fig.add_subplot(122, projection='3d')

round_probs = []
for r, round_diffs in enumerate(differential_path):
    round_prob = 1.0
    for in_diff, out_diff in round_diffs:
        round_prob *= PROB_TABLE[in_diff, out_diff]
    round_probs.append(round_prob)

x_rounds = np.arange(len(round_probs))
y_sboxes = np.arange(num_sboxes)
xpos, ypos = np.meshgrid(x_rounds, y_sboxes)
xpos = xpos.flatten()
ypos = ypos.flatten()
zpos = np.zeros_like(xpos)

dx = dy = 0.7
dz = []
colors = []

for r in range(len(differential_path)):
    for s in range(num_sboxes):
        in_diff, out_diff = differential_path[r][s]
        prob = PROB_TABLE[in_diff, out_diff]
        dz.append(prob)
        colors.append(plt.cm.plasma(prob))

ax2.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors, alpha=0.8, edgecolor='black')
ax2.set_xlabel('Round', fontsize=10)
ax2.set_ylabel('S-box', fontsize=10)
ax2.set_zlabel('Probability', fontsize=10)
ax2.set_title('Per-Sbox Differential Probabilities', fontsize=12, fontweight='bold')
ax2.view_init(elev=20, azim=135)

plt.tight_layout()
plt.savefig('3d_analysis.png', dpi=150, bbox_inches='tight')
plt.show()

# Summary statistics
print("\n" + "="*60)
print("SUMMARY STATISTICS")
print("="*60)
print(f"Number of rounds: {len(differential_path)}")
print(f"S-boxes per round: {len(differential_path[0])}")
print(f"Overall path probability: {max_probability:.6e}")
print(f"Probability in log scale: 2^({-total_cost:.4f})")

round_probs = []
for r, round_diffs in enumerate(differential_path):
    round_prob = 1.0
    for in_diff, out_diff in round_diffs:
        round_prob *= PROB_TABLE[in_diff, out_diff]
    round_probs.append(round_prob)
    print(f"Round {r} probability: {round_prob:.6e}")

print("\n" + "="*60)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import hashlib
from collections import defaultdict
import time

# Simple hash function for demonstration (truncated SHA256)
def simple_hash(x, n_bits=24):
    """Hash function that outputs n_bits"""
    h = hashlib.sha256(str(x).encode()).hexdigest()
    # Convert to integer and take lower n_bits
    return int(h, 16) % (2 ** n_bits)

# Strategy 1: Brute Force Attack (minimal memory)
def brute_force_attack(n_bits, max_attempts=1000000):
    """Try random inputs until finding a collision"""
    seen = {}
    attempts = 0
    
    for i in range(max_attempts):
        h = simple_hash(i, n_bits)
        if h in seen:
            return attempts + 1, 2  # Found collision, memory = 2 entries
        seen[h] = i
        attempts += 1
        
        if len(seen) > 100000:  # Limit memory
            seen.clear()
    
    return attempts, 2

# Strategy 2: Birthday Attack (hash table)
def birthday_attack(n_bits, table_size=None):
    """Store hashes in table until collision found"""
    if table_size is None:
        table_size = int(2 ** (n_bits / 2))
    
    hash_table = {}
    attempts = 0
    
    for i in range(table_size):
        h = simple_hash(i, n_bits)
        if h in hash_table:
            return attempts + 1, len(hash_table)
        hash_table[h] = i
        attempts += 1
    
    return attempts, len(hash_table)

# Strategy 3: Time-Memory Tradeoff
def tradeoff_attack(n_bits, memory_limit):
    """Use limited memory with multiple passes"""
    chains_per_table = memory_limit
    num_tables = max(1, int(2 ** (n_bits / 2) / memory_limit))
    
    total_attempts = 0
    
    # Build phase
    for table_idx in range(num_tables):
        hash_table = {}
        for i in range(chains_per_table):
            x = table_idx * chains_per_table + i
            h = simple_hash(x, n_bits)
            if h in hash_table:
                return total_attempts + i + 1, len(hash_table)
            hash_table[h] = x
            total_attempts += 1
    
    return total_attempts, chains_per_table

# Calculate costs for different strategies
def analyze_strategies(n_bits, alpha=1.0, beta=0.5):
    """Analyze different attack strategies"""
    results = []
    
    # Strategy 1: Brute Force
    print("Running Brute Force Attack...")
    time_bf, mem_bf = brute_force_attack(n_bits, max_attempts=50000)
    cost_bf = alpha * time_bf + beta * mem_bf
    results.append(("Brute Force", time_bf, mem_bf, cost_bf))
    
    # Strategy 2: Birthday Attack with various table sizes
    print("Running Birthday Attacks with different table sizes...")
    for table_fraction in [0.5, 0.75, 1.0, 1.25, 1.5]:
        table_size = int((2 ** (n_bits / 2)) * table_fraction)
        time_ba, mem_ba = birthday_attack(n_bits, table_size)
        cost_ba = alpha * time_ba + beta * mem_ba
        results.append((f"Birthday ({table_fraction}x)", time_ba, mem_ba, cost_ba))
    
    # Strategy 3: Tradeoff with various memory limits
    print("Running Tradeoff Attacks with different memory limits...")
    for mem_fraction in [0.1, 0.25, 0.5, 0.75, 1.0]:
        memory_limit = int((2 ** (n_bits / 2)) * mem_fraction)
        time_to, mem_to = tradeoff_attack(n_bits, memory_limit)
        cost_to = alpha * time_to + beta * mem_to
        results.append((f"Tradeoff ({mem_fraction}x)", time_to, mem_to, cost_to))
    
    return results

# Theoretical analysis
def theoretical_analysis(n_bits, alpha=1.0, beta=0.5):
    """Calculate theoretical optimal tradeoff"""
    n = n_bits
    
    # Grid of memory values (as exponent of 2)
    memory_exps = np.linspace(0, n/2, 50)
    memories = 2 ** memory_exps
    
    # For each memory, calculate time and cost
    times = []
    costs = []
    
    for M_exp in memory_exps:
        M = 2 ** M_exp
        # Time-memory tradeoff: T * M^2 = 2^n (approximately)
        T = (2 ** n) / (M ** 2)
        cost = alpha * T + beta * M
        
        times.append(T)
        costs.append(cost)
    
    return memories, np.array(times), np.array(costs)

# Main analysis
n_bits = 24
alpha = 1.0
beta = 0.5

print(f"Hash Function Collision Search Analysis (n={n_bits} bits)")
print(f"Cost parameters: α={alpha} (time), β={beta} (memory)")
print("=" * 60)

# Run practical simulations
results = analyze_strategies(n_bits, alpha, beta)

# Display results
print("\nSimulation Results:")
print("-" * 60)
print(f"{'Strategy':<25} {'Time':<12} {'Memory':<12} {'Cost':<12}")
print("-" * 60)
for name, t, m, c in results:
    print(f"{name:<25} {t:<12.0f} {m:<12.0f} {c:<12.2f}")

# Find optimal strategy
optimal = min(results, key=lambda x: x[3])
print("\n" + "=" * 60)
print(f"Optimal Strategy: {optimal[0]}")
print(f"Time: {optimal[1]:.0f}, Memory: {optimal[2]:.0f}, Cost: {optimal[3]:.2f}")
print("=" * 60)

# Theoretical analysis
memories_th, times_th, costs_th = theoretical_analysis(n_bits, alpha, beta)
optimal_idx = np.argmin(costs_th)
optimal_memory_th = memories_th[optimal_idx]
optimal_time_th = times_th[optimal_idx]
optimal_cost_th = costs_th[optimal_idx]

print(f"\nTheoretical Optimal Point:")
print(f"Memory: {optimal_memory_th:.2f}, Time: {optimal_time_th:.2f}, Cost: {optimal_cost_th:.2f}")

# Visualization
fig = plt.figure(figsize=(18, 12))

# Plot 1: Cost vs Memory (log scale)
ax1 = plt.subplot(2, 3, 1)
memories_sim = [r[2] for r in results]
costs_sim = [r[3] for r in results]
ax1.scatter(memories_sim, costs_sim, c='red', s=100, alpha=0.6, label='Simulated', zorder=5)
ax1.plot(memories_th, costs_th, 'b-', linewidth=2, label='Theoretical')
ax1.scatter([optimal_memory_th], [optimal_cost_th], c='green', s=200, marker='*', 
            label='Theoretical Optimum', zorder=10)
ax1.set_xlabel('Memory (entries)', fontsize=12)
ax1.set_ylabel('Total Cost', fontsize=12)
ax1.set_title('Cost vs Memory', fontsize=14, fontweight='bold')
ax1.set_xscale('log')
ax1.set_yscale('log')
ax1.legend()
ax1.grid(True, alpha=0.3)

# Plot 2: Time vs Memory tradeoff
ax2 = plt.subplot(2, 3, 2)
times_sim = [r[1] for r in results]
ax2.scatter(memories_sim, times_sim, c='red', s=100, alpha=0.6, label='Simulated', zorder=5)
ax2.plot(memories_th, times_th, 'b-', linewidth=2, label='Theoretical T·M²=2ⁿ')
ax2.set_xlabel('Memory (entries)', fontsize=12)
ax2.set_ylabel('Time (operations)', fontsize=12)
ax2.set_title('Time-Memory Tradeoff', fontsize=14, fontweight='bold')
ax2.set_xscale('log')
ax2.set_yscale('log')
ax2.legend()
ax2.grid(True, alpha=0.3)

# Plot 3: Cost breakdown
ax3 = plt.subplot(2, 3, 3)
strategies = [r[0] for r in results]
time_costs = [alpha * r[1] for r in results]
mem_costs = [beta * r[2] for r in results]
x_pos = np.arange(len(strategies))
ax3.bar(x_pos, time_costs, label='Time Cost', alpha=0.7)
ax3.bar(x_pos, mem_costs, bottom=time_costs, label='Memory Cost', alpha=0.7)
ax3.set_xticks(x_pos)
ax3.set_xticklabels(strategies, rotation=45, ha='right', fontsize=8)
ax3.set_ylabel('Cost', fontsize=12)
ax3.set_title('Cost Breakdown by Strategy', fontsize=14, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3, axis='y')

# Plot 4: 3D Surface - Cost as function of Time and Memory
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
M_range = np.logspace(0, n_bits/2, 30)
T_range = np.logspace(0, n_bits, 30)
M_grid, T_grid = np.meshgrid(M_range, T_range)
C_grid = alpha * T_grid + beta * M_grid

surf = ax4.plot_surface(np.log10(M_grid), np.log10(T_grid), np.log10(C_grid), 
                        cmap='viridis', alpha=0.7, edgecolor='none')
ax4.scatter([np.log10(r[2]) for r in results], 
           [np.log10(r[1]) for r in results],
           [np.log10(r[3]) for r in results],
           c='red', s=100, marker='o', label='Simulated')
ax4.set_xlabel('log₁₀(Memory)', fontsize=10)
ax4.set_ylabel('log₁₀(Time)', fontsize=10)
ax4.set_zlabel('log₁₀(Cost)', fontsize=10)
ax4.set_title('3D Cost Surface', fontsize=14, fontweight='bold')
fig.colorbar(surf, ax=ax4, shrink=0.5)

# Plot 5: 3D Tradeoff Curve with Constraint
ax5 = fig.add_subplot(2, 3, 5, projection='3d')
# Constraint surface: T * M^2 = 2^n
M_constraint = np.logspace(0, n_bits/2, 50)
T_constraint = (2 ** n_bits) / (M_constraint ** 2)
C_constraint = alpha * T_constraint + beta * M_constraint

ax5.plot(np.log10(M_constraint), np.log10(T_constraint), np.log10(C_constraint),
         'b-', linewidth=3, label='Optimal Tradeoff Curve')
ax5.scatter([np.log10(optimal_memory_th)], 
           [np.log10(optimal_time_th)],
           [np.log10(optimal_cost_th)],
           c='green', s=300, marker='*', label='Optimum')
ax5.scatter([np.log10(r[2]) for r in results], 
           [np.log10(r[1]) for r in results],
           [np.log10(r[3]) for r in results],
           c='red', s=100, marker='o', alpha=0.6, label='Simulated')
ax5.set_xlabel('log₁₀(Memory)', fontsize=10)
ax5.set_ylabel('log₁₀(Time)', fontsize=10)
ax5.set_zlabel('log₁₀(Cost)', fontsize=10)
ax5.set_title('Optimal Tradeoff with T·M²=2ⁿ Constraint', fontsize=14, fontweight='bold')
ax5.legend()

# Plot 6: Efficiency comparison
ax6 = plt.subplot(2, 3, 6)
efficiency = [optimal[3] / r[3] for r in results]
colors = ['green' if e >= 0.9 else 'orange' if e >= 0.7 else 'red' for e in efficiency]
ax6.barh(strategies, efficiency, color=colors, alpha=0.7)
ax6.axvline(x=1.0, color='black', linestyle='--', linewidth=2, label='Optimal')
ax6.set_xlabel('Efficiency (Optimal Cost / Strategy Cost)', fontsize=12)
ax6.set_title('Strategy Efficiency Comparison', fontsize=14, fontweight='bold')
ax6.set_xlim(0, 1.1)
ax6.legend()
ax6.grid(True, alpha=0.3, axis='x')

plt.tight_layout()
plt.savefig('collision_search_analysis.png', dpi=300, bbox_inches='tight')
plt.show()

print("\n" + "=" * 60)
print("Analysis complete. Visualization saved as 'collision_search_analysis.png'")
print("=" * 60)