import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.linalg import sqrtm
from scipy.optimize import differential_evolution
import seaborn as sns

# Set random seed for reproducibility
np.random.seed(42)
sns.set_style("whitegrid")

class EntanglementOptimizer:
    def __init__(self, spectrum, dim_A=2, dim_B=2):
        """
        Initialize optimizer for finding maximally entangled states.
        
        Parameters:
        - spectrum: fixed eigenvalues (must sum to 1)
        - dim_A, dim_B: dimensions of subsystems A and B
        """
        self.spectrum = np.array(sorted(spectrum, reverse=True))
        self.dim_A = dim_A
        self.dim_B = dim_B
        self.dim_total = dim_A * dim_B
        
        # Normalize spectrum
        self.spectrum = self.spectrum / np.sum(self.spectrum)
        
    def generate_density_matrix(self, params):
        """
        Generate density matrix from parameters using spectral decomposition.
        params defines a unitary matrix via Euler angles.
        """
        # Create unitary matrix from parameters
        U = self._params_to_unitary(params)
        
        # Construct density matrix: rho = U * diag(spectrum) * U^dagger
        D = np.diag(self.spectrum)
        rho = U @ D @ U.conj().T
        
        return rho
    
    def _params_to_unitary(self, params):
        """
        Convert parameter vector to unitary matrix using Givens rotations.
        """
        n = self.dim_total
        num_params_needed = n * (n - 1) // 2
        
        # Pad or truncate params
        if len(params) < num_params_needed:
            params = np.concatenate([params, np.zeros(num_params_needed - len(params))])
        else:
            params = params[:num_params_needed]
        
        U = np.eye(n, dtype=complex)
        param_idx = 0
        
        # Apply Givens rotations
        for i in range(n):
            for j in range(i + 1, n):
                theta = params[param_idx]
                # Create Givens rotation matrix
                G = np.eye(n, dtype=complex)
                G[i, i] = np.cos(theta)
                G[j, j] = np.cos(theta)
                G[i, j] = -np.sin(theta)
                G[j, i] = np.sin(theta)
                U = G @ U
                param_idx += 1
        
        return U
    
    def partial_transpose(self, rho):
        """
        Compute partial transpose with respect to subsystem B.
        """
        rho_reshaped = rho.reshape(self.dim_A, self.dim_B, self.dim_A, self.dim_B)
        rho_pt = np.transpose(rho_reshaped, (0, 3, 2, 1))
        return rho_pt.reshape(self.dim_total, self.dim_total)
    
    def negativity(self, rho):
        """
        Calculate negativity as entanglement measure.
        """
        rho_pt = self.partial_transpose(rho)
        eigenvalues = np.linalg.eigvalsh(rho_pt)
        trace_norm = np.sum(np.abs(eigenvalues))
        return (trace_norm - 1) / 2
    
    def entropy_of_entanglement(self, rho):
        """
        Calculate entropy of entanglement (von Neumann entropy of reduced state).
        """
        # Compute reduced density matrix for subsystem A
        rho_reshaped = rho.reshape(self.dim_A, self.dim_B, self.dim_A, self.dim_B)
        rho_A = np.trace(rho_reshaped, axis1=1, axis2=3)
        
        # Compute von Neumann entropy
        eigenvalues = np.linalg.eigvalsh(rho_A)
        eigenvalues = eigenvalues[eigenvalues > 1e-10]  # Remove numerical zeros
        entropy = -np.sum(eigenvalues * np.log2(eigenvalues))
        
        return entropy
    
    def objective_negativity(self, params):
        """
        Objective function: negative negativity (for minimization).
        """
        rho = self.generate_density_matrix(params)
        return -self.negativity(rho)
    
    def objective_entropy(self, params):
        """
        Objective function: negative entropy (for minimization).
        """
        rho = self.generate_density_matrix(params)
        return -self.entropy_of_entanglement(rho)
    
    def optimize(self, measure='negativity', maxiter=300):
        """
        Find optimal unitary that maximizes entanglement.
        """
        num_params = self.dim_total * (self.dim_total - 1) // 2
        bounds = [(0, 2 * np.pi)] * num_params
        
        if measure == 'negativity':
            objective = self.objective_negativity
        else:
            objective = self.objective_entropy
        
        result = differential_evolution(
            objective,
            bounds,
            maxiter=maxiter,
            seed=42,
            polish=True,
            atol=1e-8,
            tol=1e-8
        )
        
        optimal_rho = self.generate_density_matrix(result.x)
        optimal_value = -result.fun
        
        return optimal_rho, optimal_value, result

# Define test spectrum
spectrum_1 = [0.4, 0.3, 0.2, 0.1]

print("="*60)
print("MAXIMALLY ENTANGLED STATE SEARCH WITH FIXED SPECTRUM")
print("="*60)
print(f"\nFixed spectrum: {spectrum_1}")
print(f"Sum of eigenvalues: {sum(spectrum_1):.6f}")

# Create optimizer
optimizer = EntanglementOptimizer(spectrum_1, dim_A=2, dim_B=2)

# Optimize for maximum negativity
print("\n" + "-"*60)
print("Optimizing for Maximum Negativity...")
print("-"*60)
rho_optimal_neg, max_negativity, result_neg = optimizer.optimize(measure='negativity', maxiter=400)

print(f"\nMaximum Negativity Found: {max_negativity:.6f}")
print(f"Optimization Success: {result_neg.success}")
print(f"Number of iterations: {result_neg.nit}")

# Optimize for maximum entropy
print("\n" + "-"*60)
print("Optimizing for Maximum Entropy of Entanglement...")
print("-"*60)
rho_optimal_ent, max_entropy, result_ent = optimizer.optimize(measure='entropy', maxiter=400)

print(f"\nMaximum Entropy Found: {max_entropy:.6f}")
print(f"Optimization Success: {result_ent.success}")
print(f"Number of iterations: {result_ent.nit}")

# Verify spectrum preservation
print("\n" + "-"*60)
print("Verification of Spectrum Preservation")
print("-"*60)
eigenvalues_neg = np.sort(np.real(np.linalg.eigvals(rho_optimal_neg)))[::-1]
eigenvalues_ent = np.sort(np.real(np.linalg.eigvals(rho_optimal_ent)))[::-1]

print("\nOriginal Spectrum:", spectrum_1)
print("Negativity-optimized eigenvalues:", np.round(eigenvalues_neg, 6))
print("Entropy-optimized eigenvalues:", np.round(eigenvalues_ent, 6))

# Create comparison with random states
print("\n" + "-"*60)
print("Comparison with Random States")
print("-"*60)

num_random = 100
random_negativities = []
random_entropies = []

for _ in range(num_random):
    random_params = np.random.uniform(0, 2*np.pi, optimizer.dim_total * (optimizer.dim_total - 1) // 2)
    rho_random = optimizer.generate_density_matrix(random_params)
    random_negativities.append(optimizer.negativity(rho_random))
    random_entropies.append(optimizer.entropy_of_entanglement(rho_random))

print(f"\nRandom states statistics:")
print(f"Negativity - Mean: {np.mean(random_negativities):.6f}, Std: {np.std(random_negativities):.6f}")
print(f"Negativity - Min: {np.min(random_negativities):.6f}, Max: {np.max(random_negativities):.6f}")
print(f"Entropy - Mean: {np.mean(random_entropies):.6f}, Std: {np.std(random_entropies):.6f}")
print(f"Entropy - Min: {np.min(random_entropies):.6f}, Max: {np.max(random_entropies):.6f}")

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

# 1. Density matrices heatmaps
ax1 = fig.add_subplot(3, 3, 1)
im1 = ax1.imshow(np.abs(rho_optimal_neg), cmap='viridis', aspect='auto')
ax1.set_title('Optimal Density Matrix (Negativity)\n|ρ|', fontsize=11, fontweight='bold')
ax1.set_xlabel('Column Index')
ax1.set_ylabel('Row Index')
plt.colorbar(im1, ax=ax1)

ax2 = fig.add_subplot(3, 3, 2)
im2 = ax2.imshow(np.abs(rho_optimal_ent), cmap='viridis', aspect='auto')
ax2.set_title('Optimal Density Matrix (Entropy)\n|ρ|', fontsize=11, fontweight='bold')
ax2.set_xlabel('Column Index')
ax2.set_ylabel('Row Index')
plt.colorbar(im2, ax=ax2)

# 2. Partial transpose eigenvalues
ax3 = fig.add_subplot(3, 3, 3)
rho_pt_neg = optimizer.partial_transpose(rho_optimal_neg)
pt_eigenvalues = np.sort(np.real(np.linalg.eigvals(rho_pt_neg)))
ax3.bar(range(len(pt_eigenvalues)), pt_eigenvalues, color='coral', alpha=0.7, edgecolor='black')
ax3.axhline(y=0, color='red', linestyle='--', linewidth=2, label='Zero line')
ax3.set_title('Partial Transpose Eigenvalues\n(Negativity-optimized)', fontsize=11, fontweight='bold')
ax3.set_xlabel('Eigenvalue Index')
ax3.set_ylabel('Eigenvalue')
ax3.legend()
ax3.grid(True, alpha=0.3)

# 3. Comparison histogram
ax4 = fig.add_subplot(3, 3, 4)
ax4.hist(random_negativities, bins=30, alpha=0.6, color='blue', edgecolor='black', label='Random States')
ax4.axvline(max_negativity, color='red', linestyle='--', linewidth=2, label=f'Optimal: {max_negativity:.4f}')
ax4.set_title('Negativity Distribution', fontsize=11, fontweight='bold')
ax4.set_xlabel('Negativity')
ax4.set_ylabel('Frequency')
ax4.legend()
ax4.grid(True, alpha=0.3)

ax5 = fig.add_subplot(3, 3, 5)
ax5.hist(random_entropies, bins=30, alpha=0.6, color='green', edgecolor='black', label='Random States')
ax5.axvline(max_entropy, color='red', linestyle='--', linewidth=2, label=f'Optimal: {max_entropy:.4f}')
ax5.set_title('Entropy Distribution', fontsize=11, fontweight='bold')
ax5.set_xlabel('Entropy of Entanglement')
ax5.set_ylabel('Frequency')
ax5.legend()
ax5.grid(True, alpha=0.3)

# 4. Spectrum comparison
ax6 = fig.add_subplot(3, 3, 6)
x_pos = np.arange(len(spectrum_1))
width = 0.25
ax6.bar(x_pos - width, spectrum_1, width, label='Original', alpha=0.8, edgecolor='black')
ax6.bar(x_pos, eigenvalues_neg, width, label='Negativity-opt', alpha=0.8, edgecolor='black')
ax6.bar(x_pos + width, eigenvalues_ent, width, label='Entropy-opt', alpha=0.8, edgecolor='black')
ax6.set_title('Spectrum Verification', fontsize=11, fontweight='bold')
ax6.set_xlabel('Eigenvalue Index')
ax6.set_ylabel('Eigenvalue')
ax6.legend()
ax6.grid(True, alpha=0.3)

# 5. 3D scatter plot: Negativity vs Entropy vs Purity
ax7 = fig.add_subplot(3, 3, 7, projection='3d')
purities = [np.real(np.trace(optimizer.generate_density_matrix(np.random.uniform(0, 2*np.pi, 
            optimizer.dim_total * (optimizer.dim_total - 1) // 2)) @ 
            optimizer.generate_density_matrix(np.random.uniform(0, 2*np.pi, 
            optimizer.dim_total * (optimizer.dim_total - 1) // 2)))) for _ in range(num_random)]

ax7.scatter(random_negativities, random_entropies, purities, 
           c=random_negativities, cmap='coolwarm', alpha=0.5, s=20)
ax7.scatter([max_negativity], [optimizer.entropy_of_entanglement(rho_optimal_neg)],
           [np.real(np.trace(rho_optimal_neg @ rho_optimal_neg))],
           c='red', s=200, marker='*', edgecolor='black', linewidth=2, label='Optimal (Neg)')
ax7.scatter([optimizer.negativity(rho_optimal_ent)], [max_entropy],
           [np.real(np.trace(rho_optimal_ent @ rho_optimal_ent))],
           c='yellow', s=200, marker='*', edgecolor='black', linewidth=2, label='Optimal (Ent)')
ax7.set_xlabel('Negativity', fontsize=9)
ax7.set_ylabel('Entropy', fontsize=9)
ax7.set_zlabel('Purity', fontsize=9)
ax7.set_title('3D Entanglement Landscape', fontsize=11, fontweight='bold')
ax7.legend()

# 6. Real and imaginary parts of optimal density matrix
ax8 = fig.add_subplot(3, 3, 8)
im8 = ax8.imshow(np.real(rho_optimal_neg), cmap='RdBu_r', aspect='auto', 
                 vmin=-np.max(np.abs(rho_optimal_neg)), vmax=np.max(np.abs(rho_optimal_neg)))
ax8.set_title('Real Part of ρ (Negativity-opt)', fontsize=11, fontweight='bold')
ax8.set_xlabel('Column Index')
ax8.set_ylabel('Row Index')
plt.colorbar(im8, ax=ax8)

ax9 = fig.add_subplot(3, 3, 9)
im9 = ax9.imshow(np.imag(rho_optimal_neg), cmap='RdBu_r', aspect='auto',
                 vmin=-np.max(np.abs(rho_optimal_neg)), vmax=np.max(np.abs(rho_optimal_neg)))
ax9.set_title('Imaginary Part of ρ (Negativity-opt)', fontsize=11, fontweight='bold')
ax9.set_xlabel('Column Index')
ax9.set_ylabel('Row Index')
plt.colorbar(im9, ax=ax9)

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

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

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize, differential_evolution
import warnings
warnings.filterwarnings('ignore')

# Define Pauli matrices
sigma_y = np.array([[0, -1j], [1j, 0]], dtype=complex)

def compute_concurrence(rho):
    """
    Compute concurrence for a two-qubit density matrix.
    
    Parameters:
    rho: 4x4 complex density matrix
    
    Returns:
    concurrence value (float)
    """
    # Compute spin-flipped density matrix
    sigma_y_tensor = np.kron(sigma_y, sigma_y)
    rho_tilde = sigma_y_tensor @ np.conjugate(rho) @ sigma_y_tensor
    
    # Compute R = rho * rho_tilde
    R = rho @ rho_tilde
    
    # Get eigenvalues
    eigenvalues = np.linalg.eigvalsh(R)
    
    # Sort in decreasing order and take square roots
    lambdas = np.sqrt(np.maximum(eigenvalues, 0))
    lambdas = np.sort(lambdas)[::-1]
    
    # Compute concurrence
    C = max(0, lambdas[0] - lambdas[1] - lambdas[2] - lambdas[3])
    
    return C

def binary_entropy(x):
    """Binary entropy function"""
    if x <= 0 or x >= 1:
        return 0
    return -x * np.log2(x) - (1-x) * np.log2(1-x)

def entanglement_of_formation(C):
    """
    Compute entanglement of formation from concurrence.
    
    Parameters:
    C: concurrence value
    
    Returns:
    entanglement of formation (in ebits)
    """
    if C <= 0:
        return 0
    arg = (1 + np.sqrt(1 - C**2)) / 2
    return binary_entropy(arg)

def create_density_matrix(p, theta, phi):
    """
    Create a density matrix parameterized by p, theta, phi.
    
    rho = p |psi><psi| + (1-p) I/4
    where |psi> = cos(theta)|00> + e^(i*phi)sin(theta)|11>
    
    Parameters:
    p: mixing parameter (0 to 1)
    theta: angle parameter (0 to pi/2)
    phi: phase parameter (0 to 2*pi)
    
    Returns:
    4x4 density matrix
    """
    # Create pure state |psi>
    psi = np.array([
        np.cos(theta),
        0,
        0,
        np.exp(1j * phi) * np.sin(theta)
    ], dtype=complex)
    
    # Pure state density matrix
    rho_pure = np.outer(psi, np.conj(psi))
    
    # Maximally mixed state
    rho_mixed = np.eye(4, dtype=complex) / 4
    
    # Combined density matrix
    rho = p * rho_pure + (1 - p) * rho_mixed
    
    return rho

def objective_function(params):
    """
    Objective function for optimization (negative concurrence for minimization).
    
    Parameters:
    params: [p, theta, phi]
    
    Returns:
    negative concurrence
    """
    p, theta, phi = params
    
    # Enforce constraints
    p = np.clip(p, 0, 1)
    theta = np.clip(theta, 0, np.pi/2)
    phi = np.clip(phi, 0, 2*np.pi)
    
    rho = create_density_matrix(p, theta, phi)
    C = compute_concurrence(rho)
    
    return -C  # Negative for maximization

# Optimization using differential evolution (global optimization)
print("="*60)
print("OPTIMIZATION: Maximizing Concurrence")
print("="*60)

# Bounds for parameters: p in [0,1], theta in [0, pi/2], phi in [0, 2*pi]
bounds = [(0, 1), (0, np.pi/2), (0, 2*np.pi)]

result = differential_evolution(objective_function, bounds, seed=42, maxiter=1000)

optimal_p, optimal_theta, optimal_phi = result.x
optimal_concurrence = -result.fun

print(f"\nOptimal parameters:")
print(f"  p (mixing parameter) = {optimal_p:.6f}")
print(f"  θ (theta) = {optimal_theta:.6f} rad = {np.degrees(optimal_theta):.2f}°")
print(f"  φ (phi) = {optimal_phi:.6f} rad = {np.degrees(optimal_phi):.2f}°")
print(f"\nMaximum concurrence: {optimal_concurrence:.6f}")

optimal_rho = create_density_matrix(optimal_p, optimal_theta, optimal_phi)
optimal_Ef = entanglement_of_formation(optimal_concurrence)
print(f"Entanglement of formation: {optimal_Ef:.6f} ebits")

# Verify the optimal state
print("\nOptimal density matrix:")
print(optimal_rho)
print(f"\nTrace of density matrix: {np.trace(optimal_rho):.6f}")
print(f"Is Hermitian: {np.allclose(optimal_rho, optimal_rho.conj().T)}")

eigenvalues_rho = np.linalg.eigvalsh(optimal_rho)
print(f"Eigenvalues: {eigenvalues_rho}")
print(f"All eigenvalues non-negative: {np.all(eigenvalues_rho >= -1e-10)}")

print("\n" + "="*60)
print("VISUALIZATION: Parameter Space Exploration")
print("="*60)

# Create grids for visualization
n_points = 30
p_range = np.linspace(0, 1, n_points)
theta_range = np.linspace(0, np.pi/2, n_points)
phi_fixed = optimal_phi  # Fix phi at optimal value

# Compute concurrence over the grid
C_grid = np.zeros((n_points, n_points))
Ef_grid = np.zeros((n_points, n_points))

for i, p in enumerate(p_range):
    for j, theta in enumerate(theta_range):
        rho = create_density_matrix(p, theta, phi_fixed)
        C = compute_concurrence(rho)
        C_grid[i, j] = C
        Ef_grid[i, j] = entanglement_of_formation(C)

# Create figure with multiple subplots
fig = plt.figure(figsize=(18, 12))

# 1. 3D surface plot of Concurrence
ax1 = fig.add_subplot(2, 3, 1, projection='3d')
P, THETA = np.meshgrid(p_range, theta_range)
surf1 = ax1.plot_surface(P, THETA, C_grid.T, cmap='viridis', alpha=0.9, edgecolor='none')
ax1.scatter([optimal_p], [optimal_theta], [optimal_concurrence], 
            color='red', s=100, marker='*', label='Optimal')
ax1.set_xlabel('p (mixing parameter)', fontsize=10)
ax1.set_ylabel('θ (theta) [rad]', fontsize=10)
ax1.set_zlabel('Concurrence', fontsize=10)
ax1.set_title('Concurrence in (p, θ) space', fontsize=12, fontweight='bold')
ax1.legend()
fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=5)

# 2. 3D surface plot of Entanglement of Formation
ax2 = fig.add_subplot(2, 3, 2, projection='3d')
surf2 = ax2.plot_surface(P, THETA, Ef_grid.T, cmap='plasma', alpha=0.9, edgecolor='none')
ax2.scatter([optimal_p], [optimal_theta], [optimal_Ef], 
            color='red', s=100, marker='*', label='Optimal')
ax2.set_xlabel('p (mixing parameter)', fontsize=10)
ax2.set_ylabel('θ (theta) [rad]', fontsize=10)
ax2.set_zlabel('E_f (ebits)', fontsize=10)
ax2.set_title('Entanglement of Formation in (p, θ) space', fontsize=12, fontweight='bold')
ax2.legend()
fig.colorbar(surf2, ax=ax2, shrink=0.5, aspect=5)

# 3. Contour plot of Concurrence
ax3 = fig.add_subplot(2, 3, 3)
contour = ax3.contourf(P, THETA, C_grid.T, levels=20, cmap='viridis')
ax3.contour(P, THETA, C_grid.T, levels=10, colors='black', alpha=0.3, linewidths=0.5)
ax3.scatter([optimal_p], [optimal_theta], color='red', s=100, marker='*', 
            label=f'Optimal: C={optimal_concurrence:.4f}', zorder=5)
ax3.set_xlabel('p (mixing parameter)', fontsize=10)
ax3.set_ylabel('θ (theta) [rad]', fontsize=10)
ax3.set_title('Concurrence Contour Plot', fontsize=12, fontweight='bold')
ax3.legend()
ax3.grid(True, alpha=0.3)
fig.colorbar(contour, ax=ax3)

# 4. Concurrence vs p (at optimal theta)
ax4 = fig.add_subplot(2, 3, 4)
C_vs_p = [compute_concurrence(create_density_matrix(p, optimal_theta, optimal_phi)) 
          for p in p_range]
ax4.plot(p_range, C_vs_p, 'b-', linewidth=2, label='C(p)')
ax4.axvline(optimal_p, color='r', linestyle='--', alpha=0.7, label=f'Optimal p={optimal_p:.4f}')
ax4.axhline(optimal_concurrence, color='g', linestyle='--', alpha=0.7, 
            label=f'Max C={optimal_concurrence:.4f}')
ax4.set_xlabel('p (mixing parameter)', fontsize=10)
ax4.set_ylabel('Concurrence', fontsize=10)
ax4.set_title(f'Concurrence vs p (θ={optimal_theta:.4f}, φ={optimal_phi:.4f})', 
              fontsize=12, fontweight='bold')
ax4.grid(True, alpha=0.3)
ax4.legend()

# 5. Concurrence vs theta (at optimal p)
ax5 = fig.add_subplot(2, 3, 5)
C_vs_theta = [compute_concurrence(create_density_matrix(optimal_p, theta, optimal_phi)) 
              for theta in theta_range]
ax5.plot(theta_range, C_vs_theta, 'b-', linewidth=2, label='C(θ)')
ax5.axvline(optimal_theta, color='r', linestyle='--', alpha=0.7, 
            label=f'Optimal θ={optimal_theta:.4f}')
ax5.axhline(optimal_concurrence, color='g', linestyle='--', alpha=0.7, 
            label=f'Max C={optimal_concurrence:.4f}')
ax5.set_xlabel('θ (theta) [rad]', fontsize=10)
ax5.set_ylabel('Concurrence', fontsize=10)
ax5.set_title(f'Concurrence vs θ (p={optimal_p:.4f}, φ={optimal_phi:.4f})', 
              fontsize=12, fontweight='bold')
ax5.grid(True, alpha=0.3)
ax5.legend()

# 6. Relationship between Concurrence and Entanglement of Formation
ax6 = fig.add_subplot(2, 3, 6)
C_samples = np.linspace(0, 1, 100)
Ef_samples = [entanglement_of_formation(c) for c in C_samples]
ax6.plot(C_samples, Ef_samples, 'b-', linewidth=2, label='E_f(C)')
ax6.scatter([optimal_concurrence], [optimal_Ef], color='red', s=100, marker='*', 
            label=f'Optimal: C={optimal_concurrence:.4f}, E_f={optimal_Ef:.4f}', zorder=5)
ax6.set_xlabel('Concurrence', fontsize=10)
ax6.set_ylabel('Entanglement of Formation (ebits)', fontsize=10)
ax6.set_title('E_f vs Concurrence Relationship', fontsize=12, fontweight='bold')
ax6.grid(True, alpha=0.3)
ax6.legend()

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

print("\nVisualization complete! Graph saved as 'entanglement_optimization.png'")

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import differential_evolution
import pandas as pd
from matplotlib import cm

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

# Problem parameters
ALPHA = 0.6  # Weight for area
BETA = 0.4   # Weight for power
T_MIN = 1000  # Minimum throughput (Mbps)
S_MIN = 80    # Minimum security score

# Design parameter options
N_SBOX_OPTIONS = np.array([1, 2, 4, 8, 16])
N_PIPELINE_OPTIONS = np.array([1, 2, 3, 4, 5])
FREQ_MIN = 50   # MHz
FREQ_MAX = 500  # MHz

# Area model: A(n_s, n_p) = base_area + s-box_area * n_s + pipeline_overhead * n_p
def calculate_area(n_s, n_p):
    """Calculate circuit area in mm²"""
    base_area = 0.5  # Base area for control logic
    sbox_area = 0.08  # Area per S-box
    pipeline_overhead = 0.15  # Area per pipeline stage
    area = base_area + sbox_area * n_s + pipeline_overhead * n_p
    return area

# Power model: P(n_s, n_p, f) = static_power + dynamic_power(n_s, n_p) * f
def calculate_power(n_s, n_p, f):
    """Calculate power consumption in mW"""
    static_power = 5.0  # Static power
    dynamic_coeff = 0.002 * n_s + 0.001 * n_p  # Dynamic power coefficient
    power = static_power + dynamic_coeff * f
    return power

# Throughput model: T(n_s, n_p, f) considers parallelism and pipelining
def calculate_throughput(n_s, n_p, f):
    """Calculate throughput in Mbps"""
    # AES block size is 128 bits
    block_size = 128
    # Cycles per block depends on parallelism and pipelining
    cycles_per_block = max(10 / n_s, 1) / n_p
    # Throughput = (f * 10^6) * block_size / cycles_per_block / 10^6
    throughput = f * block_size / cycles_per_block
    return throughput

# Security model: S(n_s, n_p) - more S-boxes and stages improve side-channel resistance
def calculate_security(n_s, n_p):
    """Calculate security score (0-100)"""
    # More S-boxes provide better parallelism and reduce correlation
    sbox_contribution = 30 * np.log2(n_s + 1) / np.log2(17)
    # More pipeline stages reduce timing side-channels
    pipeline_contribution = 25 * (n_p - 1) / 4
    # Base security from AES algorithm
    base_security = 45
    security = base_security + sbox_contribution + pipeline_contribution
    return min(security, 100)

# Objective function
def objective_function(x):
    """
    x[0]: index for n_s (0-4)
    x[1]: index for n_p (0-4)
    x[2]: frequency (50-500 MHz)
    """
    n_s_idx = int(round(x[0]))
    n_p_idx = int(round(x[1]))
    f = x[2]
    
    # Clamp indices
    n_s_idx = np.clip(n_s_idx, 0, len(N_SBOX_OPTIONS) - 1)
    n_p_idx = np.clip(n_p_idx, 0, len(N_PIPELINE_OPTIONS) - 1)
    
    n_s = N_SBOX_OPTIONS[n_s_idx]
    n_p = N_PIPELINE_OPTIONS[n_p_idx]
    
    # Calculate metrics
    area = calculate_area(n_s, n_p)
    power = calculate_power(n_s, n_p, f)
    throughput = calculate_throughput(n_s, n_p, f)
    security = calculate_security(n_s, n_p)
    
    # Penalty for constraint violations
    penalty = 0
    if throughput < T_MIN:
        penalty += 1000 * (T_MIN - throughput)
    if security < S_MIN:
        penalty += 1000 * (S_MIN - security)
    
    # Objective: minimize weighted sum of area and power
    objective = ALPHA * area + BETA * power + penalty
    
    return objective

# Optimization using differential evolution
print("Starting optimization...")
print(f"Constraints: Throughput >= {T_MIN} Mbps, Security >= {S_MIN}")
print(f"Objective: Minimize {ALPHA}*Area + {BETA}*Power\n")

bounds = [(0, len(N_SBOX_OPTIONS) - 1), 
          (0, len(N_PIPELINE_OPTIONS) - 1), 
          (FREQ_MIN, FREQ_MAX)]

result = differential_evolution(objective_function, bounds, 
                                seed=42, maxiter=1000, 
                                popsize=30, atol=1e-6, tol=1e-6)

# Extract optimal solution
optimal_n_s_idx = int(round(result.x[0]))
optimal_n_p_idx = int(round(result.x[1]))
optimal_f = result.x[2]

optimal_n_s = N_SBOX_OPTIONS[optimal_n_s_idx]
optimal_n_p = N_PIPELINE_OPTIONS[optimal_n_p_idx]

optimal_area = calculate_area(optimal_n_s, optimal_n_p)
optimal_power = calculate_power(optimal_n_s, optimal_n_p, optimal_f)
optimal_throughput = calculate_throughput(optimal_n_s, optimal_n_p, optimal_f)
optimal_security = calculate_security(optimal_n_s, optimal_n_p)

print("=" * 60)
print("OPTIMIZATION RESULTS")
print("=" * 60)
print(f"Optimal Number of S-boxes: {optimal_n_s}")
print(f"Optimal Pipeline Stages: {optimal_n_p}")
print(f"Optimal Clock Frequency: {optimal_f:.2f} MHz")
print(f"\nPerformance Metrics:")
print(f"  Circuit Area: {optimal_area:.4f} mm²")
print(f"  Power Consumption: {optimal_power:.4f} mW")
print(f"  Throughput: {optimal_throughput:.2f} Mbps")
print(f"  Security Score: {optimal_security:.2f}")
print(f"\nObjective Value: {result.fun:.4f}")
print("=" * 60)

# Generate comprehensive data for all combinations
print("\nGenerating design space exploration data...")

results_data = []
for n_s in N_SBOX_OPTIONS:
    for n_p in N_PIPELINE_OPTIONS:
        for f in np.linspace(FREQ_MIN, FREQ_MAX, 20):
            area = calculate_area(n_s, n_p)
            power = calculate_power(n_s, n_p, f)
            throughput = calculate_throughput(n_s, n_p, f)
            security = calculate_security(n_s, n_p)
            
            feasible = (throughput >= T_MIN) and (security >= S_MIN)
            objective = ALPHA * area + BETA * power if feasible else np.nan
            
            results_data.append({
                'n_s': n_s,
                'n_p': n_p,
                'freq': f,
                'area': area,
                'power': power,
                'throughput': throughput,
                'security': security,
                'feasible': feasible,
                'objective': objective
            })

df = pd.DataFrame(results_data)
df_feasible = df[df['feasible']]

print(f"Total design points: {len(df)}")
print(f"Feasible design points: {len(df_feasible)}")

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

# Plot 1: Area vs Power (3D with frequency)
ax1 = fig.add_subplot(2, 3, 1, projection='3d')
scatter1 = ax1.scatter(df_feasible['area'], df_feasible['power'], df_feasible['freq'],
                       c=df_feasible['objective'], cmap='viridis', s=20, alpha=0.6)
ax1.scatter([optimal_area], [optimal_power], [optimal_f], 
            c='red', s=200, marker='*', edgecolors='black', linewidths=2,
            label='Optimal Solution')
ax1.set_xlabel('Area (mm²)', fontsize=10)
ax1.set_ylabel('Power (mW)', fontsize=10)
ax1.set_zlabel('Frequency (MHz)', fontsize=10)
ax1.set_title('Design Space: Area vs Power vs Frequency', fontsize=12, fontweight='bold')
cbar1 = plt.colorbar(scatter1, ax=ax1, pad=0.1, shrink=0.6)
cbar1.set_label('Objective Value', fontsize=9)
ax1.legend(fontsize=8)

# Plot 2: Throughput vs Security (colored by objective)
ax2 = fig.add_subplot(2, 3, 2)
scatter2 = ax2.scatter(df_feasible['throughput'], df_feasible['security'],
                       c=df_feasible['objective'], cmap='plasma', s=30, alpha=0.6)
ax2.scatter([optimal_throughput], [optimal_security], 
            c='red', s=300, marker='*', edgecolors='black', linewidths=2,
            label='Optimal Solution', zorder=5)
ax2.axhline(y=S_MIN, color='green', linestyle='--', linewidth=2, label=f'Min Security = {S_MIN}')
ax2.axvline(x=T_MIN, color='blue', linestyle='--', linewidth=2, label=f'Min Throughput = {T_MIN}')
ax2.set_xlabel('Throughput (Mbps)', fontsize=11)
ax2.set_ylabel('Security Score', fontsize=11)
ax2.set_title('Throughput vs Security Trade-off', fontsize=12, fontweight='bold')
ax2.legend(fontsize=9)
ax2.grid(True, alpha=0.3)
cbar2 = plt.colorbar(scatter2, ax=ax2)
cbar2.set_label('Objective Value', fontsize=9)

# Plot 3: Pareto front (Area vs Power)
ax3 = fig.add_subplot(2, 3, 3)
for n_s in N_SBOX_OPTIONS:
    df_ns = df_feasible[df_feasible['n_s'] == n_s]
    if len(df_ns) > 0:
        ax3.scatter(df_ns['area'], df_ns['power'], label=f'{n_s} S-boxes', 
                   s=40, alpha=0.6)
ax3.scatter([optimal_area], [optimal_power], 
            c='red', s=300, marker='*', edgecolors='black', linewidths=2,
            label='Optimal', zorder=5)
ax3.set_xlabel('Circuit Area (mm²)', fontsize=11)
ax3.set_ylabel('Power Consumption (mW)', fontsize=11)
ax3.set_title('Pareto Front: Area vs Power', fontsize=12, fontweight='bold')
ax3.legend(fontsize=8, loc='upper left')
ax3.grid(True, alpha=0.3)

# Plot 4: 3D surface - Throughput vs n_s and n_p
ax4 = fig.add_subplot(2, 3, 4, projection='3d')
n_s_grid = []
n_p_grid = []
throughput_grid = []
for n_s in N_SBOX_OPTIONS:
    for n_p in N_PIPELINE_OPTIONS:
        n_s_grid.append(n_s)
        n_p_grid.append(n_p)
        throughput_grid.append(calculate_throughput(n_s, n_p, 300))  # At 300 MHz

ax4.plot_trisurf(n_s_grid, n_p_grid, throughput_grid, cmap='coolwarm', alpha=0.8)
ax4.scatter([optimal_n_s], [optimal_n_p], [optimal_throughput],
            c='red', s=200, marker='*', edgecolors='black', linewidths=2)
ax4.set_xlabel('Number of S-boxes', fontsize=10)
ax4.set_ylabel('Pipeline Stages', fontsize=10)
ax4.set_zlabel('Throughput (Mbps)', fontsize=10)
ax4.set_title('Throughput Surface (f=300MHz)', fontsize=12, fontweight='bold')

# Plot 5: Security heatmap
ax5 = fig.add_subplot(2, 3, 5)
security_matrix = np.zeros((len(N_PIPELINE_OPTIONS), len(N_SBOX_OPTIONS)))
for i, n_p in enumerate(N_PIPELINE_OPTIONS):
    for j, n_s in enumerate(N_SBOX_OPTIONS):
        security_matrix[i, j] = calculate_security(n_s, n_p)

im = ax5.imshow(security_matrix, cmap='RdYlGn', aspect='auto', origin='lower')
ax5.set_xticks(range(len(N_SBOX_OPTIONS)))
ax5.set_yticks(range(len(N_PIPELINE_OPTIONS)))
ax5.set_xticklabels(N_SBOX_OPTIONS)
ax5.set_yticklabels(N_PIPELINE_OPTIONS)
ax5.set_xlabel('Number of S-boxes', fontsize=11)
ax5.set_ylabel('Pipeline Stages', fontsize=11)
ax5.set_title('Security Score Heatmap', fontsize=12, fontweight='bold')
plt.colorbar(im, ax=ax5, label='Security Score')

# Add text annotations
for i in range(len(N_PIPELINE_OPTIONS)):
    for j in range(len(N_SBOX_OPTIONS)):
        text = ax5.text(j, i, f'{security_matrix[i, j]:.1f}',
                       ha="center", va="center", color="black", fontsize=8)

# Mark optimal point
optimal_i = np.where(N_PIPELINE_OPTIONS == optimal_n_p)[0][0]
optimal_j = np.where(N_SBOX_OPTIONS == optimal_n_s)[0][0]
ax5.plot(optimal_j, optimal_i, 'r*', markersize=20, markeredgecolor='black', markeredgewidth=2)

# Plot 6: Objective value distribution
ax6 = fig.add_subplot(2, 3, 6)
objective_values = df_feasible['objective'].dropna()
ax6.hist(objective_values, bins=30, color='skyblue', edgecolor='black', alpha=0.7)
ax6.axvline(x=result.fun, color='red', linestyle='--', linewidth=3, 
            label=f'Optimal = {result.fun:.4f}')
ax6.set_xlabel('Objective Value', fontsize=11)
ax6.set_ylabel('Frequency', fontsize=11)
ax6.set_title('Distribution of Objective Values', fontsize=12, fontweight='bold')
ax6.legend(fontsize=10)
ax6.grid(True, alpha=0.3)

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

# Summary statistics
print("\n" + "=" * 60)
print("DESIGN SPACE STATISTICS")
print("=" * 60)
print(f"Area range: {df_feasible['area'].min():.4f} - {df_feasible['area'].max():.4f} mm²")
print(f"Power range: {df_feasible['power'].min():.4f} - {df_feasible['power'].max():.4f} mW")
print(f"Throughput range: {df_feasible['throughput'].min():.2f} - {df_feasible['throughput'].max():.2f} Mbps")
print(f"Security range: {df_feasible['security'].min():.2f} - {df_feasible['security'].max():.2f}")
print(f"Objective range: {df_feasible['objective'].min():.4f} - {df_feasible['objective'].max():.4f}")
print("=" * 60)

# Comparison table
print("\n" + "=" * 60)
print("COMPARISON: TOP 5 DESIGNS")
print("=" * 60)
top_designs = df_feasible.nsmallest(5, 'objective')[['n_s', 'n_p', 'freq', 'area', 'power', 'throughput', 'security', 'objective']]
print(top_designs.to_string(index=False))
print("=" * 60)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import linprog, milp, LinearConstraint, Bounds
import pandas as pd
from itertools import product

# Problem setup: Side-channel masking optimization
# We have 4 types of masked operations with different characteristics
operations = {
    'Type A (Order 1)': {'random_cost': 2, 'security': 10, 'masking_order': 1},
    'Type B (Order 2)': {'random_cost': 5, 'security': 25, 'masking_order': 2},
    'Type C (Order 3)': {'random_cost': 9, 'security': 45, 'masking_order': 3},
    'Type D (Order 4)': {'random_cost': 14, 'security': 70, 'masking_order': 4}
}

# Security requirement
required_security = 150
min_masking_order = 2  # Each operation must provide at least order-2 masking

print("=" * 70)
print("SIDE-CHANNEL COUNTERMEASURE OPTIMIZATION")
print("Minimizing Random Number Consumption with Masking Order Constraints")
print("=" * 70)
print()

# Extract parameters
op_names = list(operations.keys())
n_ops = len(op_names)
random_costs = np.array([operations[op]['random_cost'] for op in op_names])
security_values = np.array([operations[op]['security'] for op in op_names])
masking_orders = np.array([operations[op]['masking_order'] for op in op_names])

print("Operation Parameters:")
print("-" * 70)
df_ops = pd.DataFrame({
    'Operation': op_names,
    'Random Cost (per op)': random_costs,
    'Security Value': security_values,
    'Masking Order': masking_orders
})
print(df_ops.to_string(index=False))
print()
print(f"Required Total Security: {required_security}")
print(f"Minimum Masking Order per Operation: {min_masking_order}")
print()

# Linear Programming formulation
# Objective: minimize sum of (random_cost * x_i)
c = random_costs

# Inequality constraints: -security_values * x >= -required_security
# (equivalent to security_values * x >= required_security)
A_ub = -security_values.reshape(1, -1)
b_ub = np.array([-required_security])

# Masking order constraints: masking_order * x_i >= min_masking_order
# For each operation i: d_i * x_i >= D_min
# Rewrite as: -d_i * x_i <= -D_min
A_order = np.zeros((n_ops, n_ops))
b_order = np.zeros(n_ops)
for i in range(n_ops):
    if masking_orders[i] > 0:
        A_order[i, i] = -masking_orders[i]
        b_order[i] = -min_masking_order

# Combine constraints
A_ub_combined = np.vstack([A_ub, A_order])
b_ub_combined = np.hstack([b_ub, b_order])

# Bounds: x_i >= 0
bounds = [(0, None) for _ in range(n_ops)]

print("=" * 70)
print("SOLVING LINEAR PROGRAMMING RELAXATION")
print("=" * 70)
print()

# Solve LP relaxation (allows fractional solutions)
result_lp = linprog(c, A_ub=A_ub_combined, b_ub=b_ub_combined, 
                    bounds=bounds, method='highs')

if result_lp.success:
    print("LP Relaxation Solution (Fractional):")
    print("-" * 70)
    for i, op in enumerate(op_names):
        print(f"{op}: {result_lp.x[i]:.4f} operations")
    print()
    print(f"Minimum Random Numbers (Fractional): {result_lp.fun:.4f}")
    print(f"Total Security Achieved: {np.dot(security_values, result_lp.x):.4f}")
    print()
else:
    print("LP relaxation failed to find a solution")
    print()

# Integer Linear Programming (actual solution needed)
print("=" * 70)
print("SOLVING INTEGER LINEAR PROGRAMMING")
print("=" * 70)
print()

# Use scipy.optimize.milp for integer programming
from scipy.optimize import milp

# Convert to milp format
constraints_milp = LinearConstraint(A_ub_combined, -np.inf, b_ub_combined)
integrality = np.ones(n_ops)  # All variables are integers
bounds_milp = Bounds(lb=np.zeros(n_ops), ub=np.full(n_ops, 100))

result_milp = milp(c=c, constraints=constraints_milp, 
                   integrality=integrality, bounds=bounds_milp)

optimal_solution = None
optimal_cost = None

if result_milp.success:
    print("Integer Programming Solution:")
    print("-" * 70)
    optimal_solution = np.round(result_milp.x).astype(int)
    optimal_cost = np.dot(random_costs, optimal_solution)
    
    for i, op in enumerate(op_names):
        print(f"{op}: {optimal_solution[i]} operations")
    print()
    print(f"Minimum Random Numbers Required: {optimal_cost}")
    print(f"Total Security Achieved: {np.dot(security_values, optimal_solution)}")
    
    # Verify constraints
    print()
    print("Constraint Verification:")
    print("-" * 70)
    total_sec = np.dot(security_values, optimal_solution)
    print(f"Security constraint: {total_sec} >= {required_security} : {'✓' if total_sec >= required_security else '✗'}")
    
    for i in range(n_ops):
        order_achieved = masking_orders[i] * optimal_solution[i]
        print(f"Masking order {op_names[i]}: {order_achieved} >= {min_masking_order} : {'✓' if order_achieved >= min_masking_order else '✗'}")
    print()
else:
    print("Integer programming failed. Using brute force search...")
    
    # Brute force for small problem
    min_cost = float('inf')
    best_solution = None
    
    # Search reasonable range
    max_each = 20
    for combo in product(range(max_each), repeat=n_ops):
        x = np.array(combo)
        
        # Check security constraint
        if np.dot(security_values, x) < required_security:
            continue
        
        # Check masking order constraints
        valid = True
        for i in range(n_ops):
            if masking_orders[i] * x[i] < min_masking_order:
                valid = False
                break
        
        if not valid:
            continue
        
        # Calculate cost
        cost = np.dot(random_costs, x)
        if cost < min_cost:
            min_cost = cost
            best_solution = x
    
    if best_solution is not None:
        optimal_solution = best_solution
        optimal_cost = min_cost
        
        print("Brute Force Solution:")
        print("-" * 70)
        for i, op in enumerate(op_names):
            print(f"{op}: {optimal_solution[i]} operations")
        print()
        print(f"Minimum Random Numbers Required: {optimal_cost}")
        print(f"Total Security Achieved: {np.dot(security_values, optimal_solution)}")
        print()

# Comparison analysis
print("=" * 70)
print("EFFICIENCY ANALYSIS")
print("=" * 70)
print()

if optimal_solution is not None:
    # Compare with naive approach (using only highest order)
    naive_ops = np.ceil(required_security / security_values[-1]).astype(int)
    naive_solution = np.zeros(n_ops, dtype=int)
    naive_solution[-1] = naive_ops
    naive_cost = random_costs[-1] * naive_ops
    
    print("Naive Approach (using only highest-order masking):")
    print(f"Operations: {naive_ops} × {op_names[-1]}")
    print(f"Random Numbers: {naive_cost}")
    print()
    
    savings = naive_cost - optimal_cost
    savings_pct = (savings / naive_cost) * 100
    
    print(f"Optimization Savings: {savings} random numbers ({savings_pct:.1f}% reduction)")
    print()

# Visualization
print("=" * 70)
print("GENERATING VISUALIZATIONS")
print("=" * 70)
print()

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

# Plot 1: Operation distribution
ax1 = fig.add_subplot(2, 3, 1)
if optimal_solution is not None:
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A']
    bars = ax1.bar(range(n_ops), optimal_solution, color=colors, edgecolor='black', linewidth=1.5)
    ax1.set_xlabel('Operation Type', fontsize=11, fontweight='bold')
    ax1.set_ylabel('Number of Operations', fontsize=11, fontweight='bold')
    ax1.set_title('Optimal Operation Distribution', fontsize=12, fontweight='bold')
    ax1.set_xticks(range(n_ops))
    ax1.set_xticklabels([f'Type {chr(65+i)}' for i in range(n_ops)], rotation=0)
    ax1.grid(axis='y', alpha=0.3)
    
    for i, (bar, val) in enumerate(zip(bars, optimal_solution)):
        height = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width()/2., height,
                f'{int(val)}', ha='center', va='bottom', fontweight='bold')

# Plot 2: Random number consumption breakdown
ax2 = fig.add_subplot(2, 3, 2)
if optimal_solution is not None:
    random_consumption = random_costs * optimal_solution
    explode = [0.05 if x > 0 else 0 for x in random_consumption]
    colors_pie = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A']
    
    wedges, texts, autotexts = ax2.pie(random_consumption, explode=explode, 
                                         labels=[f'Type {chr(65+i)}' for i in range(n_ops)],
                                         autopct='%1.1f%%', colors=colors_pie,
                                         startangle=90, textprops={'fontweight': 'bold'})
    ax2.set_title('Random Number Consumption by Type', fontsize=12, fontweight='bold')

# Plot 3: Cost vs Security tradeoff
ax3 = fig.add_subplot(2, 3, 3)
security_range = np.linspace(100, 300, 50)
min_costs = []

for target_sec in security_range:
    # Quick approximation using LP relaxation
    b_ub_temp = np.hstack([np.array([-target_sec]), b_order])
    result_temp = linprog(c, A_ub=A_ub_combined, b_ub=b_ub_temp, 
                          bounds=bounds, method='highs')
    if result_temp.success:
        min_costs.append(result_temp.fun)
    else:
        min_costs.append(np.nan)

ax3.plot(security_range, min_costs, linewidth=2.5, color='#2C3E50')
ax3.axvline(x=required_security, color='red', linestyle='--', linewidth=2, label='Required Security')
if optimal_solution is not None:
    ax3.plot(np.dot(security_values, optimal_solution), optimal_cost, 
             'ro', markersize=12, label='Optimal Solution', zorder=5)
ax3.set_xlabel('Total Security Level', fontsize=11, fontweight='bold')
ax3.set_ylabel('Random Numbers Required', fontsize=11, fontweight='bold')
ax3.set_title('Security vs Random Cost Tradeoff', fontsize=12, fontweight='bold')
ax3.legend(fontsize=9)
ax3.grid(True, alpha=0.3)

# Plot 4: 3D visualization (Operations mix)
ax4 = fig.add_subplot(2, 3, 4, projection='3d')

# Create a mesh for two operation types
n_points = 20
x_range = np.linspace(0, 10, n_points)
y_range = np.linspace(0, 10, n_points)
X, Y = np.meshgrid(x_range, y_range)
Z = np.zeros_like(X)

for i in range(n_points):
    for j in range(n_points):
        x1, x2 = X[i, j], Y[i, j]
        # Calculate minimum Type C and D needed
        remaining_sec = max(0, required_security - security_values[0]*x1 - security_values[1]*x2)
        
        if remaining_sec == 0:
            x3_min = max(0, np.ceil(min_masking_order / masking_orders[2]))
            x4_min = max(0, np.ceil(min_masking_order / masking_orders[3]))
            Z[i, j] = (random_costs[0]*x1 + random_costs[1]*x2 + 
                      random_costs[2]*x3_min + random_costs[3]*x4_min)
        else:
            # Estimate minimum cost for remaining security
            x3_min = max(0, np.ceil(min_masking_order / masking_orders[2]))
            x4_needed = max(0, np.ceil((remaining_sec - security_values[2]*x3_min) / security_values[3]))
            x4_min = max(x4_needed, np.ceil(min_masking_order / masking_orders[3]))
            
            Z[i, j] = (random_costs[0]*x1 + random_costs[1]*x2 + 
                      random_costs[2]*x3_min + random_costs[3]*x4_min)

surf = ax4.plot_surface(X, Y, Z, cmap='viridis', alpha=0.8, edgecolor='none')

if optimal_solution is not None:
    ax4.scatter([optimal_solution[0]], [optimal_solution[1]], 
               [optimal_cost], color='red', s=200, marker='o', 
               edgecolor='black', linewidth=2, label='Optimal', zorder=10)

ax4.set_xlabel('Type A Operations', fontsize=10, fontweight='bold')
ax4.set_ylabel('Type B Operations', fontsize=10, fontweight='bold')
ax4.set_zlabel('Total Random Cost', fontsize=10, fontweight='bold')
ax4.set_title('3D Cost Surface', fontsize=12, fontweight='bold')
fig.colorbar(surf, ax=ax4, shrink=0.5, aspect=5)

# Plot 5: Masking order contribution
ax5 = fig.add_subplot(2, 3, 5)
if optimal_solution is not None:
    masking_contribution = masking_orders * optimal_solution
    x_pos = np.arange(n_ops)
    colors_bar = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A']
    
    bars = ax5.bar(x_pos, masking_contribution, color=colors_bar, 
                   edgecolor='black', linewidth=1.5)
    ax5.axhline(y=min_masking_order, color='red', linestyle='--', 
                linewidth=2, label=f'Min Required ({min_masking_order})')
    ax5.set_xlabel('Operation Type', fontsize=11, fontweight='bold')
    ax5.set_ylabel('Total Masking Order Contribution', fontsize=11, fontweight='bold')
    ax5.set_title('Masking Order Verification', fontsize=12, fontweight='bold')
    ax5.set_xticks(x_pos)
    ax5.set_xticklabels([f'Type {chr(65+i)}' for i in range(n_ops)])
    ax5.legend(fontsize=9)
    ax5.grid(axis='y', alpha=0.3)
    
    for bar, val in zip(bars, masking_contribution):
        height = bar.get_height()
        ax5.text(bar.get_x() + bar.get_width()/2., height,
                f'{int(val)}', ha='center', va='bottom', fontweight='bold')

# Plot 6: Efficiency comparison
ax6 = fig.add_subplot(2, 3, 6)
if optimal_solution is not None:
    strategies = ['Naive\n(High Order Only)', 'Optimized\n(Mixed Strategy)']
    costs_comparison = [naive_cost, optimal_cost]
    colors_comp = ['#E74C3C', '#27AE60']
    
    bars = ax6.bar(strategies, costs_comparison, color=colors_comp, 
                   edgecolor='black', linewidth=2, width=0.6)
    ax6.set_ylabel('Random Numbers Required', fontsize=11, fontweight='bold')
    ax6.set_title('Strategy Comparison', fontsize=12, fontweight='bold')
    ax6.grid(axis='y', alpha=0.3)
    
    for bar, val in zip(bars, costs_comparison):
        height = bar.get_height()
        ax6.text(bar.get_x() + bar.get_width()/2., height,
                f'{int(val)}', ha='center', va='bottom', 
                fontweight='bold', fontsize=11)
    
    # Add savings annotation
    mid_x = 0.5
    ax6.annotate('', xy=(mid_x, optimal_cost), xytext=(mid_x, naive_cost),
                arrowprops=dict(arrowstyle='<->', color='red', lw=2))
    ax6.text(mid_x + 0.15, (naive_cost + optimal_cost) / 2,
            f'-{savings}\n({savings_pct:.1f}%)', fontsize=10, 
            fontweight='bold', color='red')

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

# Additional 3D Analysis: Security-Cost-Order relationship
fig2 = plt.figure(figsize=(14, 6))

# 3D Plot 1: Cost landscape with masking order constraint
ax_3d1 = fig2.add_subplot(1, 2, 1, projection='3d')

type_b_range = np.linspace(0, 8, 25)
type_c_range = np.linspace(0, 8, 25)
B, C = np.meshgrid(type_b_range, type_c_range)
Cost_surface = np.zeros_like(B)

for i in range(B.shape[0]):
    for j in range(B.shape[1]):
        b_ops, c_ops = B[i, j], C[i, j]
        
        # Check masking constraints
        if masking_orders[1] * b_ops < min_masking_order:
            b_ops = np.ceil(min_masking_order / masking_orders[1])
        if masking_orders[2] * c_ops < min_masking_order:
            c_ops = np.ceil(min_masking_order / masking_orders[2])
        
        # Calculate remaining security needed
        current_sec = security_values[1] * b_ops + security_values[2] * c_ops
        remaining = max(0, required_security - current_sec)
        
        # Add Type D if needed
        d_ops = max(np.ceil(remaining / security_values[3]), 
                    np.ceil(min_masking_order / masking_orders[3]))
        
        Cost_surface[i, j] = (random_costs[1] * b_ops + 
                              random_costs[2] * c_ops + 
                              random_costs[3] * d_ops)

surf1 = ax_3d1.plot_surface(B, C, Cost_surface, cmap='plasma', 
                             alpha=0.7, edgecolor='none')

if optimal_solution is not None and optimal_solution[1] > 0 and optimal_solution[2] > 0:
    ax_3d1.scatter([optimal_solution[1]], [optimal_solution[2]], 
                   [optimal_cost], color='lime', s=300, marker='*',
                   edgecolor='black', linewidth=2, label='Optimal Point', zorder=100)

ax_3d1.set_xlabel('Type B Ops (Order 2)', fontsize=10, fontweight='bold')
ax_3d1.set_ylabel('Type C Ops (Order 3)', fontsize=10, fontweight='bold')
ax_3d1.set_zlabel('Total Random Cost', fontsize=10, fontweight='bold')
ax_3d1.set_title('3D Cost Landscape: Type B vs Type C', fontsize=12, fontweight='bold')
ax_3d1.view_init(elev=25, azim=45)
fig2.colorbar(surf1, ax=ax_3d1, shrink=0.5, aspect=10)

# 3D Plot 2: Pareto frontier (Security, Cost, Max Order)
ax_3d2 = fig2.add_subplot(1, 2, 2, projection='3d')

# Generate various feasible solutions
solutions_3d = []
for total_ops in range(3, 25):
    for combo in product(range(total_ops+1), repeat=n_ops):
        if sum(combo) != total_ops:
            continue
        
        x = np.array(combo)
        
        # Check masking constraints
        valid = True
        for i in range(n_ops):
            if masking_orders[i] * x[i] < min_masking_order and x[i] > 0:
                valid = False
                break
        
        if not valid:
            continue
        
        sec = np.dot(security_values, x)
        if sec < required_security:
            continue
        
        cost = np.dot(random_costs, x)
        max_order = max([masking_orders[i] * x[i] for i in range(n_ops) if x[i] > 0], default=0)
        
        solutions_3d.append((sec, cost, max_order))

if solutions_3d:
    solutions_3d = np.array(solutions_3d)
    scatter = ax_3d2.scatter(solutions_3d[:, 0], solutions_3d[:, 1], solutions_3d[:, 2],
                             c=solutions_3d[:, 1], cmap='coolwarm', s=50, alpha=0.6,
                             edgecolor='black', linewidth=0.5)
    
    if optimal_solution is not None:
        opt_sec = np.dot(security_values, optimal_solution)
        opt_max_order = max([masking_orders[i] * optimal_solution[i] 
                             for i in range(n_ops) if optimal_solution[i] > 0])
        ax_3d2.scatter([opt_sec], [optimal_cost], [opt_max_order],
                      color='lime', s=400, marker='*', edgecolor='black',
                      linewidth=2, label='Optimal', zorder=100)
    
    ax_3d2.set_xlabel('Total Security', fontsize=10, fontweight='bold')
    ax_3d2.set_ylabel('Random Cost', fontsize=10, fontweight='bold')
    ax_3d2.set_zlabel('Max Masking Order', fontsize=10, fontweight='bold')
    ax_3d2.set_title('3D Pareto Frontier Analysis', fontsize=12, fontweight='bold')
    ax_3d2.view_init(elev=20, azim=135)
    ax_3d2.legend(fontsize=9)
    fig2.colorbar(scatter, ax=ax_3d2, shrink=0.5, aspect=10, label='Cost')

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

plt.show()

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