import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize_scalar
from matplotlib import cm

# Set random seed for reproducibility
np.random.seed(42)

# Problem parameters
n = 100  # Lattice dimension
q = 1009  # Modulus for LWE
sigma = 3.0  # Noise parameter

# BKZ parameters range
beta_min = 10
beta_max = 80
beta_range = np.arange(beta_min, beta_max + 1)

# Time budget scenarios (in seconds)
time_budgets = [1, 10, 100, 1000, 10000]

print("="*70)
print("BKZ Block Size Optimization Analysis")
print("="*70)
print(f"Lattice dimension (n): {n}")
print(f"Modulus (q): {q}")
print(f"Noise parameter (σ): {sigma}")
print(f"Block size range (β): [{beta_min}, {beta_max}]")
print("="*70)

def compute_root_hermite_factor(beta, n):
    """
    Compute the root Hermite factor achieved by BKZ-β
    Using the approximation: δ ≈ ((πβ)^(1/β) * β / (2πe))^(1/(2(β-1)))
    """
    if beta <= 1:
        return 1.0
    # Chen-Nguyen approximation
    delta = ((np.pi * beta) ** (1.0 / beta) * beta / (2 * np.pi * np.e)) ** (1.0 / (2 * (beta - 1)))
    return delta

def compute_computation_time(beta):
    """
    Estimate BKZ-β running time
    T(β) ≈ c * 2^(k*β) where k ≈ 0.18 for practical implementations
    """
    c = 0.001  # Base constant (seconds)
    k = 0.18   # Exponential factor
    return c * (2 ** (k * beta))

def compute_attack_success_probability(delta, n, q, sigma):
    """
    Estimate attack success probability based on root Hermite factor
    The attack succeeds if the shortest vector found is short enough
    Success probability increases as δ decreases (better reduction)
    """
    # Expected shortest vector length after BKZ reduction
    expected_length = delta ** n * (q ** (1/n))
    
    # Target length for successful attack (related to noise)
    target_length = sigma * np.sqrt(n)
    
    # Success probability using a sigmoid function
    # Higher probability when expected_length < target_length
    ratio = expected_length / target_length
    prob = 1.0 / (1.0 + np.exp(5 * (ratio - 1)))
    
    return prob

def compute_expected_success_rate(beta, time_budget, n, q, sigma):
    """
    Compute expected success rate considering time budget
    Expected success = (time_budget / T(β)) * P_success(β)
    This represents how many successful attacks we expect within the time budget
    """
    comp_time = compute_computation_time(beta)
    if comp_time > time_budget:
        # Can't even complete one attempt
        return 0.0
    
    delta = compute_root_hermite_factor(beta, n)
    prob_success = compute_attack_success_probability(delta, n, q, sigma)
    
    # Number of attempts possible within time budget
    num_attempts = time_budget / comp_time
    
    # Expected number of successes
    expected_successes = num_attempts * prob_success
    
    return expected_successes

# Compute metrics for all block sizes
deltas = np.array([compute_root_hermite_factor(beta, n) for beta in beta_range])
comp_times = np.array([compute_computation_time(beta) for beta in beta_range])
success_probs = np.array([compute_attack_success_probability(delta, n, q, sigma) 
                          for delta in deltas])

print("\nComputed Metrics Sample (every 10 block sizes):")
print(f"{'β':<5} {'δ':<12} {'Time (s)':<15} {'P(success)':<12}")
print("-" * 50)
for i in range(0, len(beta_range), 10):
    print(f"{beta_range[i]:<5} {deltas[i]:<12.6f} {comp_times[i]:<15.4e} {success_probs[i]:<12.6f}")

# Find optimal block size for each time budget
optimal_results = []
print("\n" + "="*70)
print("Optimization Results for Different Time Budgets")
print("="*70)

for time_budget in time_budgets:
    expected_rates = np.array([compute_expected_success_rate(beta, time_budget, n, q, sigma) 
                               for beta in beta_range])
    
    optimal_idx = np.argmax(expected_rates)
    optimal_beta = beta_range[optimal_idx]
    optimal_rate = expected_rates[optimal_idx]
    
    optimal_results.append({
        'time_budget': time_budget,
        'optimal_beta': optimal_beta,
        'expected_rate': optimal_rate,
        'computation_time': compute_computation_time(optimal_beta),
        'success_prob': success_probs[optimal_idx]
    })
    
    print(f"\nTime Budget: {time_budget} seconds")
    print(f"  Optimal β: {optimal_beta}")
    print(f"  Expected successes: {optimal_rate:.4f}")
    print(f"  Single run time: {compute_computation_time(optimal_beta):.4f} seconds")
    print(f"  Single run success probability: {success_probs[optimal_idx]:.6f}")
    print(f"  Number of possible runs: {time_budget / compute_computation_time(optimal_beta):.2f}")

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

# Plot 1: Computation Time vs Block Size
ax1 = fig.add_subplot(2, 3, 1)
ax1.semilogy(beta_range, comp_times, 'b-', linewidth=2, label='Computation Time')
ax1.set_xlabel('Block Size (β)', fontsize=12)
ax1.set_ylabel('Time (seconds, log scale)', fontsize=12)
ax1.set_title('BKZ Computation Time vs Block Size', fontsize=13, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.legend()

# Plot 2: Root Hermite Factor vs Block Size
ax2 = fig.add_subplot(2, 3, 2)
ax2.plot(beta_range, deltas, 'r-', linewidth=2, label='Root Hermite Factor δ')
ax2.set_xlabel('Block Size (β)', fontsize=12)
ax2.set_ylabel('Root Hermite Factor (δ)', fontsize=12)
ax2.set_title('Root Hermite Factor vs Block Size', fontsize=13, fontweight='bold')
ax2.grid(True, alpha=0.3)
ax2.legend()

# Plot 3: Success Probability vs Block Size
ax3 = fig.add_subplot(2, 3, 3)
ax3.plot(beta_range, success_probs, 'g-', linewidth=2, label='Success Probability')
ax3.set_xlabel('Block Size (β)', fontsize=12)
ax3.set_ylabel('Attack Success Probability', fontsize=12)
ax3.set_title('Attack Success Probability vs Block Size', fontsize=13, fontweight='bold')
ax3.grid(True, alpha=0.3)
ax3.legend()

# Plot 4: Expected Success Rate for Different Time Budgets
ax4 = fig.add_subplot(2, 3, 4)
for time_budget in time_budgets:
    expected_rates = np.array([compute_expected_success_rate(beta, time_budget, n, q, sigma) 
                               for beta in beta_range])
    ax4.plot(beta_range, expected_rates, linewidth=2, label=f'T={time_budget}s', alpha=0.8)
ax4.set_xlabel('Block Size (β)', fontsize=12)
ax4.set_ylabel('Expected Number of Successes', fontsize=12)
ax4.set_title('Expected Success Rate vs Block Size\n(Different Time Budgets)', 
              fontsize=13, fontweight='bold')
ax4.grid(True, alpha=0.3)
ax4.legend()

# Plot 5: Optimal Block Size vs Time Budget
ax5 = fig.add_subplot(2, 3, 5)
budgets = [r['time_budget'] for r in optimal_results]
optimal_betas = [r['optimal_beta'] for r in optimal_results]
ax5.semilogx(budgets, optimal_betas, 'o-', linewidth=2, markersize=8, color='purple')
ax5.set_xlabel('Time Budget (seconds, log scale)', fontsize=12)
ax5.set_ylabel('Optimal Block Size (β)', fontsize=12)
ax5.set_title('Optimal Block Size vs Time Budget', fontsize=13, fontweight='bold')
ax5.grid(True, alpha=0.3)

# Plot 6: Trade-off Summary
ax6 = fig.add_subplot(2, 3, 6)
ax6_twin = ax6.twinx()
ax6.plot(beta_range, comp_times, 'b-', linewidth=2, label='Computation Time')
ax6_twin.plot(beta_range, success_probs, 'r-', linewidth=2, label='Success Probability')
ax6.set_xlabel('Block Size (β)', fontsize=12)
ax6.set_ylabel('Time (seconds)', fontsize=12, color='b')
ax6_twin.set_ylabel('Success Probability', fontsize=12, color='r')
ax6.set_title('Time vs Success Trade-off', fontsize=13, fontweight='bold')
ax6.tick_params(axis='y', labelcolor='b')
ax6_twin.tick_params(axis='y', labelcolor='r')
ax6.set_yscale('log')
ax6.grid(True, alpha=0.3)
lines1, labels1 = ax6.get_legend_handles_labels()
lines2, labels2 = ax6_twin.get_legend_handles_labels()
ax6.legend(lines1 + lines2, labels1 + labels2, loc='center left')

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

# 3D Visualization: Expected Success Rate Surface
fig_3d = plt.figure(figsize=(16, 6))

# 3D Plot 1: Success Rate Surface (Time Budget vs Block Size)
ax_3d1 = fig_3d.add_subplot(1, 2, 1, projection='3d')

# Create mesh for 3D surface
time_budget_range = np.logspace(0, 4, 50)  # 1 to 10000 seconds
beta_mesh, time_mesh = np.meshgrid(beta_range, time_budget_range)
success_rate_mesh = np.zeros_like(beta_mesh, dtype=float)

for i in range(len(time_budget_range)):
    for j in range(len(beta_range)):
        success_rate_mesh[i, j] = compute_expected_success_rate(
            beta_range[j], time_budget_range[i], n, q, sigma
        )

surf1 = ax_3d1.plot_surface(beta_mesh, np.log10(time_mesh), success_rate_mesh,
                            cmap=cm.viridis, alpha=0.9, edgecolor='none')
ax_3d1.set_xlabel('Block Size (β)', fontsize=11, labelpad=10)
ax_3d1.set_ylabel('log₁₀(Time Budget)', fontsize=11, labelpad=10)
ax_3d1.set_zlabel('Expected Successes', fontsize=11, labelpad=10)
ax_3d1.set_title('Expected Success Rate Surface\n(Time Budget vs Block Size)', 
                 fontsize=13, fontweight='bold', pad=20)
fig_3d.colorbar(surf1, ax=ax_3d1, shrink=0.5, aspect=5)

# Mark optimal points
for result in optimal_results:
    ax_3d1.scatter([result['optimal_beta']], 
                   [np.log10(result['time_budget'])], 
                   [result['expected_rate']],
                   color='red', s=100, marker='o', edgecolors='black', linewidths=2)

# 3D Plot 2: Trade-off Landscape (Root Hermite Factor vs Time vs Success)
ax_3d2 = fig_3d.add_subplot(1, 2, 2, projection='3d')

# Create data for scatter plot
scatter_betas = beta_range[::2]  # Sample every other point for clarity
scatter_deltas = deltas[::2]
scatter_times = comp_times[::2]
scatter_probs = success_probs[::2]

scatter = ax_3d2.scatter(scatter_deltas, np.log10(scatter_times), scatter_probs,
                        c=scatter_betas, cmap=cm.plasma, s=100, alpha=0.8,
                        edgecolors='black', linewidths=0.5)
ax_3d2.set_xlabel('Root Hermite Factor (δ)', fontsize=11, labelpad=10)
ax_3d2.set_ylabel('log₁₀(Time)', fontsize=11, labelpad=10)
ax_3d2.set_zlabel('Success Probability', fontsize=11, labelpad=10)
ax_3d2.set_title('Trade-off Landscape\n(Quality vs Time vs Success)', 
                 fontsize=13, fontweight='bold', pad=20)
cbar = fig_3d.colorbar(scatter, ax=ax_3d2, shrink=0.5, aspect=5)
cbar.set_label('Block Size (β)', fontsize=10)

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

# Summary statistics
print("\n" + "="*70)
print("Summary Statistics")
print("="*70)
print(f"Minimum computation time: {comp_times.min():.6f} seconds (β={beta_range[np.argmin(comp_times)]})")
print(f"Maximum computation time: {comp_times.max():.2f} seconds (β={beta_range[np.argmax(comp_times)]})")
print(f"Minimum root Hermite factor: {deltas.min():.6f} (β={beta_range[np.argmin(deltas)]})")
print(f"Maximum root Hermite factor: {deltas.max():.6f} (β={beta_range[np.argmax(deltas)]})")
print(f"Minimum success probability: {success_probs.min():.8f} (β={beta_range[np.argmin(success_probs)]})")
print(f"Maximum success probability: {success_probs.max():.6f} (β={beta_range[np.argmax(success_probs)]})")

print("\n" + "="*70)
print("Key Insights")
print("="*70)
print("1. Computation time grows exponentially with block size: T(β) ~ 2^(0.18β)")
print("2. Root Hermite factor decreases (improves) with larger block sizes")
print("3. Attack success probability increases with better reduction quality")
print("4. Optimal block size increases logarithmically with time budget")
print("5. The trade-off is non-trivial: maximum success probability doesn't always")
print("   yield maximum expected success rate within a time budget")
print("="*70)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.special import erfc
from scipy.optimize import brentq
import warnings
warnings.filterwarnings('ignore')

# Set random seed for reproducibility
np.random.seed(42)

# LWE Parameters
n = 50  # dimension
q = 1021  # modulus (prime)
sigma = 3.2  # error standard deviation

print("=== LWE Sample Count Optimization ===")
print(f"Parameters: n={n}, q={q}, σ={sigma}")
print()

def estimate_security_bits(n, q, sigma, m):
    """
    Estimate security level in bits based on LWE parameters.
    This uses a simplified model of the primal attack.
    
    The security estimate is based on:
    - Lattice dimension: d = n + m
    - Root Hermite factor: δ = (β/d)^(1/(2β-d)) where β is BKZ block size
    - Expected shortest vector length vs required length
    """
    d = n + m  # lattice dimension
    
    # Estimate BKZ block size needed (simplified)
    # For actual attacks, this would use more sophisticated models
    log_q = np.log2(q)
    log_sigma = np.log2(sigma)
    
    # Gaussian heuristic for shortest vector
    delta = 1.0045  # root Hermite factor for moderate security
    
    # Security bits approximation
    # Based on the fact that attack cost is roughly 2^(0.292*β) for BKZ-β
    beta = min(d, int(d * 0.7))  # effective block size
    security = 0.292 * beta - 16.4 + log_sigma
    
    return max(0, security)

def attack_success_probability(n, q, sigma, m):
    """
    Model the probability of successful attack given m samples.
    
    This models the success probability based on:
    1. Information-theoretic requirements (need m > n)
    2. Statistical advantage from having more samples
    3. Error accumulation effects
    """
    if m <= n:
        # Insufficient samples - very low success probability
        return 0.01 * (m / n)
    
    # Calculate effective advantage
    # More samples = more information = higher success
    excess_samples = m - n
    
    # Model: success probability increases with sample count
    # but with diminishing returns
    log_advantage = excess_samples / (2 * sigma**2)
    
    # Apply sigmoid-like function for realistic probability
    # Success probability increases as we get more samples beyond minimum
    z = (log_advantage - 2.0) / 0.5
    prob = 1.0 / (1.0 + np.exp(-z))
    
    # Factor in dimension effects
    dim_factor = np.exp(-n / 200.0)
    prob = prob * (1.0 - dim_factor) + dim_factor * 0.5
    
    return min(0.999, max(0.001, prob))

def find_optimal_samples(n, q, sigma, target_prob):
    """
    Find the minimum number of samples needed to achieve target success probability.
    Uses binary search for efficiency.
    """
    # Start with theoretical minimum
    m_min = n + 1
    m_max = n + 500
    
    # Check if target is achievable
    if attack_success_probability(n, q, sigma, m_max) < target_prob:
        print(f"Warning: Target probability {target_prob} may not be achievable with m <= {m_max}")
        return m_max
    
    # Binary search
    while m_max - m_min > 1:
        m_mid = (m_min + m_max) // 2
        prob = attack_success_probability(n, q, sigma, m_mid)
        
        if prob < target_prob:
            m_min = m_mid
        else:
            m_max = m_mid
    
    return m_max

# Find optimal sample count for different target probabilities
target_probs = [0.50, 0.75, 0.90, 0.95, 0.99]
optimal_samples = []

print("Target Probability -> Optimal Sample Count:")
print("-" * 50)

for target_prob in target_probs:
    m_opt = find_optimal_samples(n, q, sigma, target_prob)
    optimal_samples.append(m_opt)
    actual_prob = attack_success_probability(n, q, sigma, m_opt)
    security = estimate_security_bits(n, q, sigma, m_opt)
    print(f"P_target = {target_prob:.2f} -> m = {m_opt:3d} (actual P = {actual_prob:.4f}, security ≈ {security:.1f} bits)")

print()

# Generate data for visualization
m_range = np.arange(n, n + 200, 1)
success_probs = [attack_success_probability(n, q, sigma, m) for m in m_range]
security_bits = [estimate_security_bits(n, q, sigma, m) for m in m_range]

# Create 2D plots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

# Plot 1: Success Probability vs Sample Count
ax1.plot(m_range, success_probs, 'b-', linewidth=2, label='Success Probability')
ax1.axhline(y=0.95, color='r', linestyle='--', linewidth=1.5, label='Target (0.95)')
ax1.axvline(x=optimal_samples[3], color='g', linestyle='--', linewidth=1.5, 
            label=f'Optimal m={optimal_samples[3]}')
ax1.fill_between(m_range, 0, success_probs, alpha=0.3)
ax1.set_xlabel('Number of Samples (m)', fontsize=12)
ax1.set_ylabel('Attack Success Probability', fontsize=12)
ax1.set_title('LWE Attack Success Probability vs Sample Count', fontsize=13, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.legend(fontsize=10)
ax1.set_ylim(0, 1.05)

# Plot 2: Security Level vs Sample Count
ax2.plot(m_range, security_bits, 'r-', linewidth=2)
ax2.axvline(x=optimal_samples[3], color='g', linestyle='--', linewidth=1.5,
            label=f'Optimal m={optimal_samples[3]}')
ax2.fill_between(m_range, 0, security_bits, alpha=0.3, color='red')
ax2.set_xlabel('Number of Samples (m)', fontsize=12)
ax2.set_ylabel('Security Level (bits)', fontsize=12)
ax2.set_title('Remaining Security Level vs Sample Count', fontsize=13, fontweight='bold')
ax2.grid(True, alpha=0.3)
ax2.legend(fontsize=10)

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

print("2D plots generated successfully!")
print()

# Create 3D visualization
# Vary both dimension and sample count
dimensions = np.arange(30, 80, 5)
samples_3d = np.arange(40, 150, 5)
D, M = np.meshgrid(dimensions, samples_3d)

# Calculate success probability for each combination
Z = np.zeros_like(D, dtype=float)
for i in range(D.shape[0]):
    for j in range(D.shape[1]):
        n_val = D[i, j]
        m_val = M[i, j]
        Z[i, j] = attack_success_probability(n_val, q, sigma, m_val)

# Create 3D surface plot
fig = plt.figure(figsize=(14, 10))

# 3D Surface plot
ax1 = fig.add_subplot(121, projection='3d')
surf = ax1.plot_surface(D, M, Z, cmap='viridis', alpha=0.9, 
                        edgecolor='none', antialiased=True)
ax1.set_xlabel('Dimension (n)', fontsize=11, labelpad=10)
ax1.set_ylabel('Sample Count (m)', fontsize=11, labelpad=10)
ax1.set_zlabel('Success Probability', fontsize=11, labelpad=10)
ax1.set_title('3D Surface: Attack Success Probability', fontsize=12, fontweight='bold', pad=20)
ax1.view_init(elev=25, azim=45)
fig.colorbar(surf, ax=ax1, shrink=0.5, aspect=5)

# 3D Contour plot
ax2 = fig.add_subplot(122, projection='3d')
contours = ax2.contour3D(D, M, Z, 50, cmap='plasma', alpha=0.8)
ax2.set_xlabel('Dimension (n)', fontsize=11, labelpad=10)
ax2.set_ylabel('Sample Count (m)', fontsize=11, labelpad=10)
ax2.set_zlabel('Success Probability', fontsize=11, labelpad=10)
ax2.set_title('3D Contour: Attack Success Probability', fontsize=12, fontweight='bold', pad=20)
ax2.view_init(elev=25, azim=45)
fig.colorbar(contours, ax=ax2, shrink=0.5, aspect=5)

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

print("3D plots generated successfully!")
print()

# Trade-off analysis
print("=== Trade-off Analysis ===")
print("Sample Count vs Security Level for P_target = 0.95:")
print("-" * 60)

trade_off_samples = [optimal_samples[3] - 20, optimal_samples[3], optimal_samples[3] + 20]
for m in trade_off_samples:
    if m > n:
        prob = attack_success_probability(n, q, sigma, m)
        sec = estimate_security_bits(n, q, sigma, m)
        efficiency = prob / m  # success probability per sample
        print(f"m = {m:3d}: P_success = {prob:.4f}, Security ≈ {sec:.1f} bits, Efficiency = {efficiency:.6f}")

print()
print("=== Conclusion ===")
print(f"For target success probability of 0.95:")
print(f"  - Minimum samples needed: {optimal_samples[3]}")
print(f"  - This is {optimal_samples[3] - n} samples above the theoretical minimum of n={n}")
print(f"  - Remaining security: ≈{estimate_security_bits(n, q, sigma, optimal_samples[3]):.1f} bits")
print()
print("Key insight: Adding more samples beyond the optimal point")
print("provides diminishing returns in success probability while")
print("continuously reducing the security of the system.")
print()
print("### Execution Results ###")
# Placeholder for execution results

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.linalg import qr
import time

def babai_algorithm(basis, target):
    """
    Babai's nearest plane algorithm for CVP approximation
    
    Args:
        basis: numpy array of shape (n, n) representing lattice basis
        target: numpy array of shape (n,) representing target point
    
    Returns:
        closest_point: approximation of closest lattice point
        coefficients: integer coefficients
    """
    # Gram-Schmidt orthogonalization
    Q, R = qr(basis.T)
    
    # Solve for coefficients in the orthogonal basis
    coeffs_real = np.linalg.solve(basis.T, target)
    
    # Round to nearest integers (Babai's rounding)
    coeffs_int = np.round(coeffs_real).astype(int)
    
    # Compute the closest lattice point
    closest_point = basis.T @ coeffs_int
    
    return closest_point, coeffs_int

def enumerate_cvp(basis, target, search_radius=5):
    """
    Brute force enumeration to find exact CVP solution
    
    Args:
        basis: numpy array of shape (n, n)
        target: numpy array of shape (n,)
        search_radius: how far to search in coefficient space
    
    Returns:
        best_point: exact closest lattice point
        best_coeffs: corresponding coefficients
        min_distance: minimum distance found
    """
    n = basis.shape[1]
    best_distance = float('inf')
    best_point = None
    best_coeffs = None
    
    # Generate all integer combinations within search radius
    ranges = [range(-search_radius, search_radius + 1) for _ in range(n)]
    
    import itertools
    for coeffs in itertools.product(*ranges):
        coeffs_array = np.array(coeffs)
        lattice_point = basis.T @ coeffs_array
        distance = np.linalg.norm(target - lattice_point)
        
        if distance < best_distance:
            best_distance = distance
            best_point = lattice_point
            best_coeffs = coeffs_array
    
    return best_point, best_coeffs, best_distance

def generate_lattice_points(basis, range_limit=5):
    """
    Generate lattice points for visualization
    """
    n = basis.shape[1]
    points = []
    
    import itertools
    ranges = [range(-range_limit, range_limit + 1) for _ in range(n)]
    
    for coeffs in itertools.product(*ranges):
        coeffs_array = np.array(coeffs)
        point = basis.T @ coeffs_array
        points.append(point)
    
    return np.array(points)

# Main execution
print("=" * 60)
print("Closest Vector Problem (CVP) Solver")
print("=" * 60)

# Define the 2D lattice basis
basis_2d = np.array([[3, 1],
                      [1, 2]])

target_2d = np.array([5.7, 3.2])

print("\nLattice Basis B:")
print(basis_2d)
print(f"\nTarget Point t: {target_2d}")

# Solve using Babai's algorithm
print("\n" + "=" * 60)
print("Method 1: Babai's Nearest Plane Algorithm (Fast Approximation)")
print("=" * 60)

start_time = time.time()
babai_point, babai_coeffs = babai_algorithm(basis_2d, target_2d)
babai_time = time.time() - start_time

babai_distance = np.linalg.norm(target_2d - babai_point)

print(f"Coefficients: {babai_coeffs}")
print(f"Closest Point (approx): {babai_point}")
print(f"Distance: {babai_distance:.6f}")
print(f"Computation Time: {babai_time:.6f} seconds")

# Solve using enumeration (exact solution)
print("\n" + "=" * 60)
print("Method 2: Exhaustive Enumeration (Exact Solution)")
print("=" * 60)

start_time = time.time()
exact_point, exact_coeffs, exact_distance = enumerate_cvp(basis_2d, target_2d, search_radius=10)
enum_time = time.time() - start_time

print(f"Coefficients: {exact_coeffs}")
print(f"Closest Point (exact): {exact_point}")
print(f"Distance: {exact_distance:.6f}")
print(f"Computation Time: {enum_time:.6f} seconds")

# 3D Example
print("\n" + "=" * 60)
print("3D Lattice Example")
print("=" * 60)

basis_3d = np.array([[4, 1, 0],
                      [1, 3, 1],
                      [0, 1, 3]])

target_3d = np.array([7.5, 5.2, 4.8])

print("\n3D Lattice Basis B:")
print(basis_3d)
print(f"\nTarget Point t: {target_3d}")

# Babai for 3D
print("\nBabai's Algorithm (3D):")
babai_point_3d, babai_coeffs_3d = babai_algorithm(basis_3d, target_3d)
babai_distance_3d = np.linalg.norm(target_3d - babai_point_3d)

print(f"Coefficients: {babai_coeffs_3d}")
print(f"Closest Point: {babai_point_3d}")
print(f"Distance: {babai_distance_3d:.6f}")

# Exact solution for 3D
print("\nExhaustive Enumeration (3D):")
start_time = time.time()
exact_point_3d, exact_coeffs_3d, exact_distance_3d = enumerate_cvp(basis_3d, target_3d, search_radius=5)
enum_time_3d = time.time() - start_time

print(f"Coefficients: {exact_coeffs_3d}")
print(f"Closest Point: {exact_point_3d}")
print(f"Distance: {exact_distance_3d:.6f}")
print(f"Computation Time: {enum_time_3d:.6f} seconds")

# Visualization
print("\n" + "=" * 60)
print("Generating Visualizations...")
print("=" * 60)

# Generate lattice points for 2D
lattice_points_2d = generate_lattice_points(basis_2d, range_limit=5)

# Generate lattice points for 3D
lattice_points_3d = generate_lattice_points(basis_3d, range_limit=3)

# Create comprehensive visualization
fig = plt.figure(figsize=(20, 12))

# 2D Plot - Overview
ax1 = fig.add_subplot(2, 3, 1)
ax1.scatter(lattice_points_2d[:, 0], lattice_points_2d[:, 1], 
           c='lightblue', s=50, alpha=0.6, label='Lattice Points')
ax1.scatter(target_2d[0], target_2d[1], 
           c='red', s=200, marker='*', label='Target', zorder=5)
ax1.scatter(exact_point[0], exact_point[1], 
           c='green', s=150, marker='s', label='Closest Point (Exact)', zorder=5)
ax1.scatter(babai_point[0], babai_point[1], 
           c='orange', s=150, marker='^', label='Babai Approximation', zorder=5)

# Draw basis vectors
origin = np.array([0, 0])
ax1.arrow(origin[0], origin[1], basis_2d[0, 0], basis_2d[1, 0], 
         head_width=0.3, head_length=0.3, fc='blue', ec='blue', linewidth=2, alpha=0.7)
ax1.arrow(origin[0], origin[1], basis_2d[0, 1], basis_2d[1, 1], 
         head_width=0.3, head_length=0.3, fc='purple', ec='purple', linewidth=2, alpha=0.7)

ax1.plot([target_2d[0], exact_point[0]], [target_2d[1], exact_point[1]], 
        'g--', linewidth=2, label=f'Distance: {exact_distance:.3f}')

ax1.set_xlabel('X', fontsize=12)
ax1.set_ylabel('Y', fontsize=12)
ax1.set_title('2D CVP: Lattice and Target Point', fontsize=14, fontweight='bold')
ax1.legend(fontsize=10)
ax1.grid(True, alpha=0.3)
ax1.axis('equal')

# 2D Plot - Zoomed in
ax2 = fig.add_subplot(2, 3, 2)
nearby_points = lattice_points_2d[
    (np.abs(lattice_points_2d[:, 0] - target_2d[0]) < 5) & 
    (np.abs(lattice_points_2d[:, 1] - target_2d[1]) < 5)
]
ax2.scatter(nearby_points[:, 0], nearby_points[:, 1], 
           c='lightblue', s=100, alpha=0.8, label='Nearby Lattice Points')
ax2.scatter(target_2d[0], target_2d[1], 
           c='red', s=300, marker='*', label='Target', zorder=5)
ax2.scatter(exact_point[0], exact_point[1], 
           c='green', s=200, marker='s', label='Closest Point', zorder=5)

for point in nearby_points:
    distance = np.linalg.norm(point - target_2d)
    ax2.plot([target_2d[0], point[0]], [target_2d[1], point[1]], 
            'gray', alpha=0.3, linewidth=1)
    
ax2.plot([target_2d[0], exact_point[0]], [target_2d[1], exact_point[1]], 
        'g-', linewidth=3, label=f'Min Distance: {exact_distance:.3f}')

ax2.set_xlabel('X', fontsize=12)
ax2.set_ylabel('Y', fontsize=12)
ax2.set_title('2D CVP: Zoomed View', fontsize=14, fontweight='bold')
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)
ax2.axis('equal')

# Distance comparison plot
ax3 = fig.add_subplot(2, 3, 3)
methods = ['Babai\n(Approx)', 'Enumeration\n(Exact)']
distances = [babai_distance, exact_distance]
colors = ['orange', 'green']
bars = ax3.bar(methods, distances, color=colors, alpha=0.7, edgecolor='black', linewidth=2)

for i, (bar, dist) in enumerate(zip(bars, distances)):
    height = bar.get_height()
    ax3.text(bar.get_x() + bar.get_width()/2., height,
            f'{dist:.4f}',
            ha='center', va='bottom', fontsize=12, fontweight='bold')

ax3.set_ylabel('Distance to Target', fontsize=12)
ax3.set_title('2D CVP: Method Comparison', fontsize=14, fontweight='bold')
ax3.grid(True, alpha=0.3, axis='y')

# 3D Plot - Main view
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
ax4.scatter(lattice_points_3d[:, 0], lattice_points_3d[:, 1], lattice_points_3d[:, 2],
           c='lightblue', s=30, alpha=0.4, label='Lattice Points')
ax4.scatter(target_3d[0], target_3d[1], target_3d[2],
           c='red', s=300, marker='*', label='Target', zorder=5)
ax4.scatter(exact_point_3d[0], exact_point_3d[1], exact_point_3d[2],
           c='green', s=200, marker='s', label='Closest Point', zorder=5)
ax4.plot([target_3d[0], exact_point_3d[0]], 
        [target_3d[1], exact_point_3d[1]], 
        [target_3d[2], exact_point_3d[2]],
        'g-', linewidth=3, label=f'Distance: {exact_distance_3d:.3f}')

# Draw basis vectors
origin_3d = np.array([0, 0, 0])
for i in range(3):
    ax4.quiver(origin_3d[0], origin_3d[1], origin_3d[2],
              basis_3d[0, i], basis_3d[1, i], basis_3d[2, i],
              arrow_length_ratio=0.1, linewidth=2, alpha=0.7)

ax4.set_xlabel('X', fontsize=12)
ax4.set_ylabel('Y', fontsize=12)
ax4.set_zlabel('Z', fontsize=12)
ax4.set_title('3D CVP: Lattice Structure', fontsize=14, fontweight='bold')
ax4.legend(fontsize=10)
ax4.grid(True, alpha=0.3)

# 3D Plot - Different angle
ax5 = fig.add_subplot(2, 3, 5, projection='3d')
nearby_points_3d = lattice_points_3d[
    (np.abs(lattice_points_3d[:, 0] - target_3d[0]) < 6) & 
    (np.abs(lattice_points_3d[:, 1] - target_3d[1]) < 6) &
    (np.abs(lattice_points_3d[:, 2] - target_3d[2]) < 6)
]
ax5.scatter(nearby_points_3d[:, 0], nearby_points_3d[:, 1], nearby_points_3d[:, 2],
           c='lightblue', s=60, alpha=0.6, label='Nearby Lattice Points')
ax5.scatter(target_3d[0], target_3d[1], target_3d[2],
           c='red', s=300, marker='*', label='Target', zorder=5)
ax5.scatter(exact_point_3d[0], exact_point_3d[1], exact_point_3d[2],
           c='green', s=200, marker='s', label='Closest Point', zorder=5)
ax5.plot([target_3d[0], exact_point_3d[0]], 
        [target_3d[1], exact_point_3d[1]], 
        [target_3d[2], exact_point_3d[2]],
        'g-', linewidth=3)

ax5.set_xlabel('X', fontsize=12)
ax5.set_ylabel('Y', fontsize=12)
ax5.set_zlabel('Z', fontsize=12)
ax5.set_title('3D CVP: Zoomed View', fontsize=14, fontweight='bold')
ax5.legend(fontsize=10)
ax5.view_init(elev=20, azim=45)
ax5.grid(True, alpha=0.3)

# 3D Distance comparison
ax6 = fig.add_subplot(2, 3, 6)
methods_3d = ['Babai\n(Approx)', 'Enumeration\n(Exact)']
distances_3d = [babai_distance_3d, exact_distance_3d]
colors_3d = ['orange', 'green']
bars_3d = ax6.bar(methods_3d, distances_3d, color=colors_3d, alpha=0.7, edgecolor='black', linewidth=2)

for i, (bar, dist) in enumerate(zip(bars_3d, distances_3d)):
    height = bar.get_height()
    ax6.text(bar.get_x() + bar.get_width()/2., height,
            f'{dist:.4f}',
            ha='center', va='bottom', fontsize=12, fontweight='bold')

ax6.set_ylabel('Distance to Target', fontsize=12)
ax6.set_title('3D CVP: Method Comparison', fontsize=14, fontweight='bold')
ax6.grid(True, alpha=0.3, axis='y')

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

print("\nVisualization complete!")
print("=" * 60)

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

def gram_schmidt(basis):
    """
    Gram-Schmidt orthogonalization for lattice basis
    """
    basis = np.array(basis, dtype=float)
    n = len(basis)
    ortho = np.zeros_like(basis)
    mu = np.zeros((n, n))
    
    for i in range(n):
        ortho[i] = basis[i].copy()
        for j in range(i):
            mu[i, j] = np.dot(basis[i], ortho[j]) / np.dot(ortho[j], ortho[j])
            ortho[i] -= mu[i, j] * ortho[j]
    
    return ortho, mu

def lll_reduction(basis, delta=0.75):
    """
    LLL algorithm for lattice basis reduction
    delta: reduction parameter (typically 0.75)
    """
    basis = np.array(basis, dtype=float)
    n = len(basis)
    
    k = 1
    while k < n:
        # Size reduction
        ortho, mu = gram_schmidt(basis)
        
        for j in range(k-1, -1, -1):
            if abs(mu[k, j]) > 0.5:
                basis[k] -= np.round(mu[k, j]) * basis[j]
                ortho, mu = gram_schmidt(basis)
        
        # Lovász condition
        ortho_norm_k = np.dot(ortho[k], ortho[k])
        ortho_norm_k_minus_1 = np.dot(ortho[k-1], ortho[k-1])
        
        if ortho_norm_k >= (delta - mu[k, k-1]**2) * ortho_norm_k_minus_1:
            k += 1
        else:
            basis[[k, k-1]] = basis[[k-1, k]]
            k = max(k-1, 1)
    
    return basis

def exhaustive_search_svp(basis, search_bound=5):
    """
    Exhaustive search for the shortest vector (exact solution for small lattices)
    search_bound: range of coefficients to search
    """
    basis = np.array(basis)
    min_norm = float('inf')
    shortest_vector = None
    shortest_coeffs = None
    
    # Generate all possible integer combinations
    ranges = [range(-search_bound, search_bound + 1) for _ in range(len(basis))]
    
    for coeffs in product(*ranges):
        if all(c == 0 for c in coeffs):
            continue
        
        vector = sum(c * basis[i] for i, c in enumerate(coeffs))
        norm = np.linalg.norm(vector)
        
        if norm < min_norm:
            min_norm = norm
            shortest_vector = vector
            shortest_coeffs = coeffs
    
    return shortest_vector, min_norm, shortest_coeffs

def generate_lattice_points(basis, coeffs_range=3):
    """
    Generate lattice points for visualization
    """
    basis = np.array(basis)
    points = []
    
    ranges = [range(-coeffs_range, coeffs_range + 1) for _ in range(len(basis))]
    
    for coeffs in product(*ranges):
        point = sum(c * basis[i] for i, c in enumerate(coeffs))
        points.append(point)
    
    return np.array(points)

# Example 1: 2D Lattice
print("=" * 60)
print("Example 1: 2D Lattice Problem")
print("=" * 60)

basis_2d = np.array([
    [4, 1],
    [1, 3]
])

print("\nOriginal Basis:")
print(basis_2d)
print("\nBasis vectors:")
print(f"b1 = {basis_2d[0]}, ||b1|| = {np.linalg.norm(basis_2d[0]):.4f}")
print(f"b2 = {basis_2d[1]}, ||b2|| = {np.linalg.norm(basis_2d[1]):.4f}")

# LLL reduction
reduced_basis_2d = lll_reduction(basis_2d.copy())
print("\nLLL-Reduced Basis:")
print(reduced_basis_2d)
print(f"b1' = {reduced_basis_2d[0]}, ||b1'|| = {np.linalg.norm(reduced_basis_2d[0]):.4f}")
print(f"b2' = {reduced_basis_2d[1]}, ||b2'|| = {np.linalg.norm(reduced_basis_2d[1]):.4f}")

# Exhaustive search
shortest_2d, min_norm_2d, coeffs_2d = exhaustive_search_svp(basis_2d, search_bound=5)
print(f"\nExact Shortest Vector (exhaustive search):")
print(f"v = {shortest_2d}")
print(f"||v|| = {min_norm_2d:.4f}")
print(f"Coefficients: {coeffs_2d}")

# Example 2: 3D Lattice
print("\n" + "=" * 60)
print("Example 2: 3D Lattice Problem")
print("=" * 60)

basis_3d = np.array([
    [5, 2, 1],
    [1, 4, 2],
    [2, 1, 5]
])

print("\nOriginal Basis:")
print(basis_3d)
for i, b in enumerate(basis_3d):
    print(f"b{i+1} = {b}, ||b{i+1}|| = {np.linalg.norm(b):.4f}")

# LLL reduction
reduced_basis_3d = lll_reduction(basis_3d.copy())
print("\nLLL-Reduced Basis:")
print(reduced_basis_3d)
for i, b in enumerate(reduced_basis_3d):
    print(f"b{i+1}' = {b}, ||b{i+1}'|| = {np.linalg.norm(b):.4f}")

# Exhaustive search
shortest_3d, min_norm_3d, coeffs_3d = exhaustive_search_svp(basis_3d, search_bound=4)
print(f"\nExact Shortest Vector (exhaustive search):")
print(f"v = {shortest_3d}")
print(f"||v|| = {min_norm_3d:.4f}")
print(f"Coefficients: {coeffs_3d}")

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

# 2D Lattice Visualization
ax1 = fig.add_subplot(2, 3, 1)
lattice_points_2d = generate_lattice_points(basis_2d, coeffs_range=3)
ax1.scatter(lattice_points_2d[:, 0], lattice_points_2d[:, 1], 
           alpha=0.3, s=20, c='gray', label='Lattice points')
ax1.quiver(0, 0, basis_2d[0, 0], basis_2d[0, 1], 
          angles='xy', scale_units='xy', scale=1, color='blue', 
          width=0.006, label='Original basis')
ax1.quiver(0, 0, basis_2d[1, 0], basis_2d[1, 1], 
          angles='xy', scale_units='xy', scale=1, color='blue', width=0.006)
ax1.quiver(0, 0, shortest_2d[0], shortest_2d[1], 
          angles='xy', scale_units='xy', scale=1, color='red', 
          width=0.008, label=f'Shortest vector (||v||={min_norm_2d:.2f})')
ax1.grid(True, alpha=0.3)
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_title('2D Lattice - Original Basis')
ax1.legend()
ax1.axis('equal')

# 2D Reduced Lattice
ax2 = fig.add_subplot(2, 3, 2)
lattice_points_2d_reduced = generate_lattice_points(reduced_basis_2d, coeffs_range=3)
ax2.scatter(lattice_points_2d_reduced[:, 0], lattice_points_2d_reduced[:, 1], 
           alpha=0.3, s=20, c='gray', label='Lattice points')
ax2.quiver(0, 0, reduced_basis_2d[0, 0], reduced_basis_2d[0, 1], 
          angles='xy', scale_units='xy', scale=1, color='green', 
          width=0.006, label='LLL-reduced basis')
ax2.quiver(0, 0, reduced_basis_2d[1, 0], reduced_basis_2d[1, 1], 
          angles='xy', scale_units='xy', scale=1, color='green', width=0.006)
ax2.quiver(0, 0, shortest_2d[0], shortest_2d[1], 
          angles='xy', scale_units='xy', scale=1, color='red', 
          width=0.008, label=f'Shortest vector')
ax2.grid(True, alpha=0.3)
ax2.set_xlabel('x')
ax2.set_ylabel('y')
ax2.set_title('2D Lattice - LLL Reduced')
ax2.legend()
ax2.axis('equal')

# Norm comparison 2D
ax3 = fig.add_subplot(2, 3, 3)
basis_names = ['b1', 'b2', "b1'", "b2'", 'shortest']
norms = [np.linalg.norm(basis_2d[0]), np.linalg.norm(basis_2d[1]),
         np.linalg.norm(reduced_basis_2d[0]), np.linalg.norm(reduced_basis_2d[1]),
         min_norm_2d]
colors_bar = ['blue', 'blue', 'green', 'green', 'red']
ax3.bar(basis_names, norms, color=colors_bar, alpha=0.7)
ax3.set_ylabel('Norm ||v||')
ax3.set_title('2D Vector Norms Comparison')
ax3.grid(True, alpha=0.3, axis='y')

# 3D Lattice Visualization
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
lattice_points_3d = generate_lattice_points(basis_3d, coeffs_range=2)
ax4.scatter(lattice_points_3d[:, 0], lattice_points_3d[:, 1], 
           lattice_points_3d[:, 2], alpha=0.2, s=15, c='gray')
for i, b in enumerate(basis_3d):
    ax4.quiver(0, 0, 0, b[0], b[1], b[2], color='blue', 
              arrow_length_ratio=0.1, linewidth=2, alpha=0.7)
ax4.quiver(0, 0, 0, shortest_3d[0], shortest_3d[1], shortest_3d[2], 
          color='red', arrow_length_ratio=0.1, linewidth=3, alpha=0.9)
ax4.set_xlabel('x')
ax4.set_ylabel('y')
ax4.set_zlabel('z')
ax4.set_title('3D Lattice - Original Basis')

# 3D Reduced Lattice
ax5 = fig.add_subplot(2, 3, 5, projection='3d')
lattice_points_3d_reduced = generate_lattice_points(reduced_basis_3d, coeffs_range=2)
ax5.scatter(lattice_points_3d_reduced[:, 0], lattice_points_3d_reduced[:, 1], 
           lattice_points_3d_reduced[:, 2], alpha=0.2, s=15, c='gray')
for i, b in enumerate(reduced_basis_3d):
    ax5.quiver(0, 0, 0, b[0], b[1], b[2], color='green', 
              arrow_length_ratio=0.1, linewidth=2, alpha=0.7)
ax5.quiver(0, 0, 0, shortest_3d[0], shortest_3d[1], shortest_3d[2], 
          color='red', arrow_length_ratio=0.1, linewidth=3, alpha=0.9)
ax5.set_xlabel('x')
ax5.set_ylabel('y')
ax5.set_zlabel('z')
ax5.set_title('3D Lattice - LLL Reduced')

# Norm comparison 3D
ax6 = fig.add_subplot(2, 3, 6)
basis_names_3d = ['b1', 'b2', 'b3', "b1'", "b2'", "b3'", 'shortest']
norms_3d = [np.linalg.norm(basis_3d[0]), np.linalg.norm(basis_3d[1]), 
           np.linalg.norm(basis_3d[2]),
           np.linalg.norm(reduced_basis_3d[0]), np.linalg.norm(reduced_basis_3d[1]),
           np.linalg.norm(reduced_basis_3d[2]), min_norm_3d]
colors_bar_3d = ['blue', 'blue', 'blue', 'green', 'green', 'green', 'red']
ax6.bar(basis_names_3d, norms_3d, color=colors_bar_3d, alpha=0.7)
ax6.set_ylabel('Norm ||v||')
ax6.set_title('3D Vector Norms Comparison')
ax6.grid(True, alpha=0.3, axis='y')
plt.xticks(rotation=45)

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

print("\n" + "=" * 60)
print("Analysis Complete!")
print("=" * 60)

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

# Simplified elliptic curve point arithmetic for secp256k1
class ECPoint:
    """Elliptic curve point for y^2 = x^3 + 7 (secp256k1)"""
    # secp256k1 parameters
    p = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
    a = 0
    b = 7
    
    def __init__(self, x, y, infinity=False):
        self.x = x
        self.y = y
        self.infinity = infinity
        
    def __add__(self, other):
        """Point addition"""
        if self.infinity:
            return other
        if other.infinity:
            return self
        
        if self.x == other.x:
            if self.y == other.y:
                return self.double()
            else:
                return ECPoint(0, 0, infinity=True)
        
        # Point addition formula
        s = ((other.y - self.y) * pow(other.x - self.x, -1, self.p)) % self.p
        x3 = (s * s - self.x - other.x) % self.p
        y3 = (s * (self.x - x3) - self.y) % self.p
        
        return ECPoint(x3, y3)
    
    def double(self):
        """Point doubling"""
        if self.infinity:
            return self
        
        s = ((3 * self.x * self.x + self.a) * pow(2 * self.y, -1, self.p)) % self.p
        x3 = (s * s - 2 * self.x) % self.p
        y3 = (s * (self.x - x3) - self.y) % self.p
        
        return ECPoint(x3, y3)
    
    def __eq__(self, other):
        if self.infinity and other.infinity:
            return True
        if self.infinity or other.infinity:
            return False
        return self.x == other.x and self.y == other.y

# secp256k1 generator point
G = ECPoint(
    0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798,
    0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8
)

class AdditionChainCounter:
    """Count operations for different scalar multiplication methods"""
    
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.doublings = 0
        self.additions = 0
    
    def total(self):
        return self.doublings + self.additions

# Method 1: Binary (Double-and-Add)
def scalar_mult_binary(k, P, counter=None):
    """Binary method for scalar multiplication"""
    if counter:
        counter.reset()
    
    if k == 0:
        return ECPoint(0, 0, infinity=True)
    
    result = ECPoint(0, 0, infinity=True)
    addend = P
    
    while k:
        if k & 1:
            result = result + addend
            if counter:
                counter.additions += 1
        addend = addend.double()
        if counter:
            counter.doublings += 1
        k >>= 1
    
    if counter:
        counter.doublings -= 1  # Last doubling not needed
    
    return result

# Method 2: NAF (Non-Adjacent Form)
def to_naf(k):
    """Convert integer to Non-Adjacent Form"""
    naf = []
    while k > 0:
        if k & 1:
            width = 2
            naf_i = k % width
            if naf_i > width // 2:
                naf_i = naf_i - width
            k = k - naf_i
        else:
            naf_i = 0
        naf.append(naf_i)
        k = k // 2
    return naf

def scalar_mult_naf(k, P, counter=None):
    """NAF method for scalar multiplication"""
    if counter:
        counter.reset()
    
    naf = to_naf(k)
    result = ECPoint(0, 0, infinity=True)
    
    for i in range(len(naf) - 1, -1, -1):
        result = result.double()
        if counter:
            counter.doublings += 1
        
        if naf[i] == 1:
            result = result + P
            if counter:
                counter.additions += 1
        elif naf[i] == -1:
            result = result + ECPoint(P.x, (-P.y) % P.p)
            if counter:
                counter.additions += 1
    
    return result

# Method 3: Window method (width=4)
def scalar_mult_window(k, P, w=4, counter=None):
    """Window method with precomputation"""
    if counter:
        counter.reset()
    
    # Precompute [P, 3P, 5P, 7P, 9P, 11P, 13P, 15P]
    precomp = {}
    precomp[1] = P
    P2 = P.double()
    
    if counter:
        counter.doublings += 1  # For computing 2P
    
    for i in range(3, 2**w, 2):
        precomp[i] = precomp[i-2] + P2
        if counter:
            counter.additions += 1
    
    # Convert k to base 2^w representation
    result = ECPoint(0, 0, infinity=True)
    bits = bin(k)[2:]
    
    i = 0
    while i < len(bits):
        result = result.double()
        if counter:
            counter.doublings += 1
        
        if bits[i] == '1':
            # Find the window
            j = min(i + w, len(bits))
            while j > i and bits[i:j].count('1') == 0:
                j -= 1
            
            window_val = int(bits[i:j], 2)
            
            # Make it odd
            while window_val % 2 == 0 and window_val > 0:
                result = result.double()
                if counter:
                    counter.doublings += 1
                window_val //= 2
                i += 1
            
            if window_val in precomp:
                result = result + precomp[window_val]
                if counter:
                    counter.additions += 1
            
            i = j
        else:
            i += 1
    
    return result

# Method 4: Optimized addition chain (for small k)
def find_addition_chain(n, max_depth=20):
    """Find short addition chain using BFS"""
    if n == 1:
        return [1]
    
    from collections import deque
    
    queue = deque([(1, [1])])
    visited = {1}
    
    while queue:
        current, chain = queue.popleft()
        
        if len(chain) > max_depth:
            break
        
        # Try all possible additions from the chain
        for i in range(len(chain)):
            for j in range(i, len(chain)):
                next_val = chain[i] + chain[j]
                
                if next_val == n:
                    return chain + [next_val]
                
                if next_val < n and next_val not in visited:
                    visited.add(next_val)
                    queue.append((next_val, chain + [next_val]))
    
    return None

def scalar_mult_optimal_chain(k, P, counter=None):
    """Scalar multiplication using optimized addition chain"""
    if counter:
        counter.reset()
    
    chain = find_addition_chain(k)
    if chain is None:
        # Fallback to binary method
        return scalar_mult_binary(k, P, counter)
    
    # Execute the chain
    values = {1: P}
    
    for i in range(1, len(chain)):
        target = chain[i]
        # Find which two previous values sum to target
        found = False
        for j in range(i):
            for l in range(j, i):
                if chain[j] + chain[l] == target:
                    if chain[j] == chain[l]:
                        values[target] = values[chain[j]].double()
                        if counter:
                            counter.doublings += 1
                    else:
                        values[target] = values[chain[j]] + values[chain[l]]
                        if counter:
                            counter.additions += 1
                    found = True
                    break
            if found:
                break
    
    return values[k]

# Comparison and analysis
def compare_methods(test_scalars):
    """Compare all methods for different scalar values"""
    results = {
        'Binary': {'ops': [], 'times': []},
        'NAF': {'ops': [], 'times': []},
        'Window': {'ops': [], 'times': []},
        'Optimal': {'ops': [], 'times': []}
    }
    
    counter = AdditionChainCounter()
    
    for k in test_scalars:
        # Binary method
        start = time.time()
        scalar_mult_binary(k, G, counter)
        results['Binary']['times'].append(time.time() - start)
        results['Binary']['ops'].append(counter.total())
        
        # NAF method
        start = time.time()
        scalar_mult_naf(k, G, counter)
        results['NAF']['times'].append(time.time() - start)
        results['NAF']['ops'].append(counter.total())
        
        # Window method
        start = time.time()
        scalar_mult_window(k, G, 4, counter)
        results['Window']['times'].append(time.time() - start)
        results['Window']['ops'].append(counter.total())
        
        # Optimal chain (only for small k)
        if k < 1000:
            start = time.time()
            scalar_mult_optimal_chain(k, G, counter)
            results['Optimal']['times'].append(time.time() - start)
            results['Optimal']['ops'].append(counter.total())
        else:
            results['Optimal']['times'].append(None)
            results['Optimal']['ops'].append(None)
    
    return results

# Test with various scalar values
print("=" * 80)
print("ECC SCALAR MULTIPLICATION - ADDITION CHAIN OPTIMIZATION ANALYSIS")
print("=" * 80)
print()

# Test scalars of increasing size
test_scalars = [15, 31, 63, 127, 255, 511, 1023, 2047, 4095]

print("Testing scalar values:", test_scalars)
print()

results = compare_methods(test_scalars)

# Display detailed results
print("OPERATION COUNTS BY METHOD:")
print("-" * 80)
print(f"{'Scalar':<10} {'Binary':<12} {'NAF':<12} {'Window':<12} {'Optimal':<12}")
print("-" * 80)

for i, k in enumerate(test_scalars):
    optimal_str = str(results['Optimal']['ops'][i]) if results['Optimal']['ops'][i] else 'N/A'
    print(f"{k:<10} {results['Binary']['ops'][i]:<12} {results['NAF']['ops'][i]:<12} "
          f"{results['Window']['ops'][i]:<12} {optimal_str:<12}")

print()
print("THEORETICAL ANALYSIS:")
print("-" * 80)

# Analyze for k=255 in detail
k_example = 255
print(f"\nDetailed analysis for k = {k_example} (binary: {bin(k_example)})")
print()

counter = AdditionChainCounter()
scalar_mult_binary(k_example, G, counter)
print(f"Binary method:")
print(f"  Doublings: {counter.doublings}")
print(f"  Additions: {counter.additions}")
print(f"  Total: {counter.total()}")
print()

scalar_mult_naf(k_example, G, counter)
print(f"NAF method:")
print(f"  Doublings: {counter.doublings}")
print(f"  Additions: {counter.additions}")
print(f"  Total: {counter.total()}")
print(f"  NAF representation: {to_naf(k_example)}")
print()

scalar_mult_window(k_example, G, 4, counter)
print(f"Window method (w=4):")
print(f"  Doublings: {counter.doublings}")
print(f"  Additions: {counter.additions}")
print(f"  Total: {counter.total()}")
print()

# For smaller example, show optimal chain
k_small = 15
chain = find_addition_chain(k_small)
print(f"Optimal addition chain for k = {k_small}:")
print(f"  Chain: {chain}")
print(f"  Length: {len(chain) - 1}")
print(f"  Binary method would need: {bin(k_small).count('1') + len(bin(k_small)) - 3} operations")
print()

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

# Plot 1: Operation counts comparison
ax1 = fig.add_subplot(2, 3, 1)
x_pos = np.arange(len(test_scalars))
width = 0.2

small_scalars = [i for i, k in enumerate(test_scalars) if k < 1000]
large_scalars = [i for i, k in enumerate(test_scalars) if k >= 1000]

ax1.bar(x_pos - 1.5*width, results['Binary']['ops'], width, label='Binary', alpha=0.8)
ax1.bar(x_pos - 0.5*width, results['NAF']['ops'], width, label='NAF', alpha=0.8)
ax1.bar(x_pos + 0.5*width, results['Window']['ops'], width, label='Window', alpha=0.8)
ax1.bar([x_pos[i] + 1.5*width for i in small_scalars], 
        [results['Optimal']['ops'][i] for i in small_scalars], 
        width, label='Optimal', alpha=0.8)

ax1.set_xlabel('Scalar Value')
ax1.set_ylabel('Total Operations')
ax1.set_title('Operation Count Comparison')
ax1.set_xticks(x_pos)
ax1.set_xticklabels(test_scalars, rotation=45)
ax1.legend()
ax1.grid(True, alpha=0.3)

# Plot 2: Efficiency ratio (relative to binary)
ax2 = fig.add_subplot(2, 3, 2)
binary_ops = results['Binary']['ops']
naf_ratio = [results['NAF']['ops'][i] / binary_ops[i] for i in range(len(test_scalars))]
window_ratio = [results['Window']['ops'][i] / binary_ops[i] for i in range(len(test_scalars))]

ax2.plot(test_scalars, naf_ratio, 'o-', label='NAF vs Binary', linewidth=2, markersize=8)
ax2.plot(test_scalars, window_ratio, 's-', label='Window vs Binary', linewidth=2, markersize=8)
ax2.axhline(y=1.0, color='r', linestyle='--', label='Binary baseline')
ax2.set_xlabel('Scalar Value')
ax2.set_ylabel('Efficiency Ratio')
ax2.set_title('Relative Efficiency (lower is better)')
ax2.set_xscale('log')
ax2.legend()
ax2.grid(True, alpha=0.3)

# Plot 3: Bit length vs operations
ax3 = fig.add_subplot(2, 3, 3)
bit_lengths = [len(bin(k)) - 2 for k in test_scalars]

ax3.plot(bit_lengths, results['Binary']['ops'], 'o-', label='Binary', linewidth=2, markersize=8)
ax3.plot(bit_lengths, results['NAF']['ops'], 's-', label='NAF', linewidth=2, markersize=8)
ax3.plot(bit_lengths, results['Window']['ops'], '^-', label='Window', linewidth=2, markersize=8)
ax3.plot(bit_lengths, bit_lengths, 'r--', label='Theoretical minimum (bit length)', linewidth=1)

ax3.set_xlabel('Bit Length of Scalar')
ax3.set_ylabel('Total Operations')
ax3.set_title('Operations vs Bit Length')
ax3.legend()
ax3.grid(True, alpha=0.3)

# Plot 4: 3D visualization of operation breakdown
ax4 = fig.add_subplot(2, 3, 4, projection='3d')

methods = ['Binary', 'NAF', 'Window']
colors = ['blue', 'green', 'orange']

for idx, method in enumerate(methods):
    doublings = []
    additions = []
    
    for i, k in enumerate(test_scalars[:6]):  # First 6 for clarity
        counter = AdditionChainCounter()
        if method == 'Binary':
            scalar_mult_binary(k, G, counter)
        elif method == 'NAF':
            scalar_mult_naf(k, G, counter)
        else:
            scalar_mult_window(k, G, 4, counter)
        
        doublings.append(counter.doublings)
        additions.append(counter.additions)
    
    ax4.plot([idx] * len(test_scalars[:6]), test_scalars[:6], doublings, 
             'o-', color=colors[idx], label=f'{method} (D)', markersize=6)
    ax4.plot([idx] * len(test_scalars[:6]), test_scalars[:6], additions, 
             's--', color=colors[idx], label=f'{method} (A)', markersize=6, alpha=0.6)

ax4.set_xlabel('Method')
ax4.set_ylabel('Scalar Value')
ax4.set_zlabel('Operation Count')
ax4.set_title('3D View: Doublings vs Additions')
ax4.set_xticks([0, 1, 2])
ax4.set_xticklabels(methods)
ax4.legend(fontsize=8)

# Plot 5: Hamming weight analysis
ax5 = fig.add_subplot(2, 3, 5)
hamming_weights = [bin(k).count('1') for k in test_scalars]

ax5.scatter(hamming_weights, results['Binary']['ops'], s=100, alpha=0.6, label='Binary')
ax5.scatter(hamming_weights, results['NAF']['ops'], s=100, alpha=0.6, label='NAF')
ax5.scatter(hamming_weights, results['Window']['ops'], s=100, alpha=0.6, label='Window')

ax5.set_xlabel('Hamming Weight (# of 1s in binary)')
ax5.set_ylabel('Total Operations')
ax5.set_title('Impact of Hamming Weight')
ax5.legend()
ax5.grid(True, alpha=0.3)

# Plot 6: Savings percentage
ax6 = fig.add_subplot(2, 3, 6)
naf_savings = [(binary_ops[i] - results['NAF']['ops'][i]) / binary_ops[i] * 100 
               for i in range(len(test_scalars))]
window_savings = [(binary_ops[i] - results['Window']['ops'][i]) / binary_ops[i] * 100 
                  for i in range(len(test_scalars))]

ax6.bar(x_pos - width/2, naf_savings, width, label='NAF', alpha=0.8)
ax6.bar(x_pos + width/2, window_savings, width, label='Window', alpha=0.8)
ax6.set_xlabel('Scalar Value')
ax6.set_ylabel('Operations Saved (%)')
ax6.set_title('Percentage Improvement over Binary')
ax6.set_xticks(x_pos)
ax6.set_xticklabels(test_scalars, rotation=45)
ax6.axhline(y=0, color='black', linestyle='-', linewidth=0.5)
ax6.legend()
ax6.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('ecc_addition_chain_analysis.png', dpi=150, bbox_inches='tight')
print("Graphs saved as 'ecc_addition_chain_analysis.png'")
print()

# Additional 3D surface plot
fig2 = plt.figure(figsize=(14, 10))

# Create data for surface plot
scalar_range = range(10, 256, 5)
methods_3d = ['Binary', 'NAF', 'Window']

ax_3d = fig2.add_subplot(111, projection='3d')

for method_idx, method in enumerate(methods_3d):
    ops_list = []
    counter = AdditionChainCounter()
    
    for k in scalar_range:
        if method == 'Binary':
            scalar_mult_binary(k, G, counter)
        elif method == 'NAF':
            scalar_mult_naf(k, G, counter)
        else:
            scalar_mult_window(k, G, 4, counter)
        ops_list.append(counter.total())
    
    ax_3d.plot(list(scalar_range), [method_idx] * len(scalar_range), ops_list, 
               label=method, linewidth=3, alpha=0.8)

ax_3d.set_xlabel('Scalar Value (k)', fontsize=12)
ax_3d.set_ylabel('Method', fontsize=12)
ax_3d.set_zlabel('Total Operations', fontsize=12)
ax_3d.set_yticks([0, 1, 2])
ax_3d.set_yticklabels(methods_3d)
ax_3d.set_title('3D Comparison: Scalar Value vs Operations by Method', fontsize=14)
ax_3d.legend(fontsize=10)
ax_3d.view_init(elev=20, azim=45)

plt.savefig('ecc_3d_surface.png', dpi=150, bbox_inches='tight')
print("3D graph saved as 'ecc_3d_surface.png'")
print()

print("=" * 80)
print("ANALYSIS COMPLETE")
print("=" * 80)