import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.integrate import odeint, solve_ivp
from scipy.optimize import minimize
import time

# Pauli matrices
sigma_x = np.array([[0, 1], [1, 0]], dtype=complex)
sigma_y = np.array([[0, -1j], [1j, 0]], dtype=complex)
sigma_z = np.array([[1, 0], [0, -1]], dtype=complex)
sigma_plus = np.array([[0, 1], [0, 0]], dtype=complex)
sigma_minus = np.array([[0, 0], [1, 0]], dtype=complex)
identity = np.eye(2, dtype=complex)

# System parameters
omega_0 = 2.0 * np.pi  # Qubit frequency (rad/s)
gamma = 0.1  # Dissipation rate
T_final = 5.0  # Final time
n_steps = 100  # Number of time steps
dt = T_final / n_steps

# Initial and target states
rho_initial = np.array([[1, 0], [0, 0]], dtype=complex)  # |0><0|
rho_target = np.array([[0, 0], [0, 1]], dtype=complex)   # |1><1|

def commutator(A, B):
    """Compute commutator [A, B]"""
    return A @ B - B @ A

def anticommutator(A, B):
    """Compute anticommutator {A, B}"""
    return A @ B + B @ A

def lindblad_rhs(rho, H, L, gamma):
    """Right-hand side of Lindblad equation"""
    coherent = -1j * commutator(H, rho)
    dissipator = gamma * (L @ rho @ L.conj().T - 0.5 * anticommutator(L.conj().T @ L, rho))
    return coherent + dissipator

def evolve_lindblad(rho0, u_values, t_values):
    """Evolve density matrix under Lindblad equation with control"""
    n_t = len(t_values)
    rho_flat = rho0.flatten()
    
    def rhs(t, rho_flat):
        rho = rho_flat.reshape((2, 2))
        # Interpolate control value
        idx = int(t / dt)
        if idx >= len(u_values):
            idx = len(u_values) - 1
        u = u_values[idx]
        
        H = 0.5 * omega_0 * sigma_z + u * sigma_x
        drho_dt = lindblad_rhs(rho, H, sigma_minus, gamma)
        return drho_dt.flatten()
    
    sol = solve_ivp(rhs, [0, T_final], rho_flat, t_eval=t_values, method='RK45')
    rho_trajectory = sol.y.T.reshape((n_t, 2, 2))
    
    return rho_trajectory

def compute_cost(u_values, t_values, rho_trajectory):
    """Compute total cost: fidelity error + control energy + dissipation"""
    # Final state fidelity error
    rho_final = rho_trajectory[-1]
    fidelity = np.real(np.trace(rho_final @ rho_target))
    fidelity_error = 1 - fidelity
    
    # Control cost
    control_cost = np.sum(u_values**2) * dt
    
    # Energy dissipation
    dissipation = 0
    for i, rho in enumerate(rho_trajectory):
        dissipation += gamma * np.real(np.trace(sigma_minus.conj().T @ sigma_minus @ rho)) * dt
    
    # Weighted total cost
    alpha_fidelity = 100.0  # Weight for fidelity
    alpha_control = 0.1     # Weight for control cost
    alpha_dissipation = 1.0 # Weight for dissipation
    
    total_cost = (alpha_fidelity * fidelity_error + 
                  alpha_control * control_cost + 
                  alpha_dissipation * dissipation)
    
    return total_cost, fidelity_error, control_cost, dissipation

def objective_function(u_flat):
    """Objective function for optimization"""
    u_values = u_flat
    t_values = np.linspace(0, T_final, n_steps)
    rho_trajectory = evolve_lindblad(rho_initial, u_values, t_values)
    total_cost, _, _, _ = compute_cost(u_values, t_values, rho_trajectory)
    return total_cost

# Optimization
print("Starting optimization...")
start_time = time.time()

# Initial guess: sinusoidal control
t_values = np.linspace(0, T_final, n_steps)
u_initial = 2.0 * np.sin(omega_0 * t_values)

# Optimize
result = minimize(objective_function, u_initial, method='L-BFGS-B', 
                  options={'maxiter': 50, 'disp': True})

u_optimal = result.x
optimization_time = time.time() - start_time

print(f"\nOptimization completed in {optimization_time:.2f} seconds")

# Compute trajectories for both initial and optimal control
rho_trajectory_initial = evolve_lindblad(rho_initial, u_initial, t_values)
rho_trajectory_optimal = evolve_lindblad(rho_initial, u_optimal, t_values)

# Compute costs
cost_initial, fid_err_initial, ctrl_cost_initial, diss_initial = compute_cost(
    u_initial, t_values, rho_trajectory_initial)
cost_optimal, fid_err_optimal, ctrl_cost_optimal, diss_optimal = compute_cost(
    u_optimal, t_values, rho_trajectory_optimal)

print(f"\nInitial Control:")
print(f"  Fidelity Error: {fid_err_initial:.6f}")
print(f"  Control Cost: {ctrl_cost_initial:.6f}")
print(f"  Dissipation: {diss_initial:.6f}")
print(f"  Total Cost: {cost_initial:.6f}")

print(f"\nOptimal Control:")
print(f"  Fidelity Error: {fid_err_optimal:.6f}")
print(f"  Control Cost: {ctrl_cost_optimal:.6f}")
print(f"  Dissipation: {diss_optimal:.6f}")
print(f"  Total Cost: {cost_optimal:.6f}")

# Extract populations and coherences
def extract_observables(rho_trajectory):
    n_t = len(rho_trajectory)
    pop_0 = np.zeros(n_t)
    pop_1 = np.zeros(n_t)
    coherence_real = np.zeros(n_t)
    coherence_imag = np.zeros(n_t)
    energy = np.zeros(n_t)
    
    for i, rho in enumerate(rho_trajectory):
        pop_0[i] = np.real(rho[0, 0])
        pop_1[i] = np.real(rho[1, 1])
        coherence_real[i] = np.real(rho[0, 1])
        coherence_imag[i] = np.imag(rho[0, 1])
        H = 0.5 * omega_0 * sigma_z
        energy[i] = np.real(np.trace(H @ rho))
    
    return pop_0, pop_1, coherence_real, coherence_imag, energy

obs_initial = extract_observables(rho_trajectory_initial)
obs_optimal = extract_observables(rho_trajectory_optimal)

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

# Plot 1: Control fields
ax1 = plt.subplot(3, 3, 1)
ax1.plot(t_values, u_initial, 'b--', label='Initial Control', linewidth=2)
ax1.plot(t_values, u_optimal, 'r-', label='Optimal Control', linewidth=2)
ax1.set_xlabel('Time (s)', fontsize=12)
ax1.set_ylabel('Control Field u(t)', fontsize=12)
ax1.set_title('Control Field Comparison', fontsize=14, fontweight='bold')
ax1.legend(fontsize=10)
ax1.grid(True, alpha=0.3)

# Plot 2: Population dynamics (initial control)
ax2 = plt.subplot(3, 3, 2)
ax2.plot(t_values, obs_initial[0], 'b-', label='P(|0⟩)', linewidth=2)
ax2.plot(t_values, obs_initial[1], 'r-', label='P(|1⟩)', linewidth=2)
ax2.set_xlabel('Time (s)', fontsize=12)
ax2.set_ylabel('Population', fontsize=12)
ax2.set_title('Initial Control: Populations', fontsize=14, fontweight='bold')
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)
ax2.set_ylim([-0.1, 1.1])

# Plot 3: Population dynamics (optimal control)
ax3 = plt.subplot(3, 3, 3)
ax3.plot(t_values, obs_optimal[0], 'b-', label='P(|0⟩)', linewidth=2)
ax3.plot(t_values, obs_optimal[1], 'r-', label='P(|1⟩)', linewidth=2)
ax3.set_xlabel('Time (s)', fontsize=12)
ax3.set_ylabel('Population', fontsize=12)
ax3.set_title('Optimal Control: Populations', fontsize=14, fontweight='bold')
ax3.legend(fontsize=10)
ax3.grid(True, alpha=0.3)
ax3.set_ylim([-0.1, 1.1])

# Plot 4: Coherence comparison
ax4 = plt.subplot(3, 3, 4)
coherence_mag_initial = np.sqrt(obs_initial[2]**2 + obs_initial[3]**2)
coherence_mag_optimal = np.sqrt(obs_optimal[2]**2 + obs_optimal[3]**2)
ax4.plot(t_values, coherence_mag_initial, 'b--', label='Initial', linewidth=2)
ax4.plot(t_values, coherence_mag_optimal, 'r-', label='Optimal', linewidth=2)
ax4.set_xlabel('Time (s)', fontsize=12)
ax4.set_ylabel('|ρ₀₁|', fontsize=12)
ax4.set_title('Coherence Magnitude', fontsize=14, fontweight='bold')
ax4.legend(fontsize=10)
ax4.grid(True, alpha=0.3)

# Plot 5: Energy dynamics
ax5 = plt.subplot(3, 3, 5)
ax5.plot(t_values, obs_initial[4], 'b--', label='Initial', linewidth=2)
ax5.plot(t_values, obs_optimal[4], 'r-', label='Optimal', linewidth=2)
ax5.set_xlabel('Time (s)', fontsize=12)
ax5.set_ylabel('Energy ⟨H⟩', fontsize=12)
ax5.set_title('System Energy', fontsize=14, fontweight='bold')
ax5.legend(fontsize=10)
ax5.grid(True, alpha=0.3)

# Plot 6: Dissipation rate
ax6 = plt.subplot(3, 3, 6)
diss_rate_initial = gamma * obs_initial[1]  # Dissipation rate proportional to excited state population
diss_rate_optimal = gamma * obs_optimal[1]
ax6.plot(t_values, diss_rate_initial, 'b--', label='Initial', linewidth=2)
ax6.plot(t_values, diss_rate_optimal, 'r-', label='Optimal', linewidth=2)
ax6.set_xlabel('Time (s)', fontsize=12)
ax6.set_ylabel('Dissipation Rate', fontsize=12)
ax6.set_title('Instantaneous Dissipation', fontsize=14, fontweight='bold')
ax6.legend(fontsize=10)
ax6.grid(True, alpha=0.3)

# Plot 7: Cost breakdown
ax7 = plt.subplot(3, 3, 7)
categories = ['Fidelity\nError', 'Control\nCost', 'Dissipation']
initial_costs = [fid_err_initial, ctrl_cost_initial, diss_initial]
optimal_costs = [fid_err_optimal, ctrl_cost_optimal, diss_optimal]
x = np.arange(len(categories))
width = 0.35
ax7.bar(x - width/2, initial_costs, width, label='Initial', color='blue', alpha=0.7)
ax7.bar(x + width/2, optimal_costs, width, label='Optimal', color='red', alpha=0.7)
ax7.set_ylabel('Cost Components', fontsize=12)
ax7.set_title('Cost Breakdown', fontsize=14, fontweight='bold')
ax7.set_xticks(x)
ax7.set_xticklabels(categories, fontsize=10)
ax7.legend(fontsize=10)
ax7.grid(True, alpha=0.3, axis='y')

# Plot 8: 3D Bloch sphere trajectory (initial)
ax8 = plt.subplot(3, 3, 8, projection='3d')
# Compute Bloch vector
bloch_x_initial = np.real([np.trace(sigma_x @ rho) for rho in rho_trajectory_initial])
bloch_y_initial = np.real([np.trace(sigma_y @ rho) for rho in rho_trajectory_initial])
bloch_z_initial = np.real([np.trace(sigma_z @ rho) for rho in rho_trajectory_initial])

# Draw Bloch sphere
u_sphere = np.linspace(0, 2 * np.pi, 30)
v_sphere = np.linspace(0, np.pi, 20)
x_sphere = np.outer(np.cos(u_sphere), np.sin(v_sphere))
y_sphere = np.outer(np.sin(u_sphere), np.sin(v_sphere))
z_sphere = np.outer(np.ones(np.size(u_sphere)), np.cos(v_sphere))
ax8.plot_surface(x_sphere, y_sphere, z_sphere, color='cyan', alpha=0.1)

# Plot trajectory
ax8.plot(bloch_x_initial, bloch_y_initial, bloch_z_initial, 'b-', linewidth=2, label='Trajectory')
ax8.scatter([bloch_x_initial[0]], [bloch_y_initial[0]], [bloch_z_initial[0]], 
            color='green', s=100, label='Start')
ax8.scatter([bloch_x_initial[-1]], [bloch_y_initial[-1]], [bloch_z_initial[-1]], 
            color='red', s=100, label='End')
ax8.set_xlabel('X', fontsize=10)
ax8.set_ylabel('Y', fontsize=10)
ax8.set_zlabel('Z', fontsize=10)
ax8.set_title('Initial Control: Bloch Sphere', fontsize=14, fontweight='bold')
ax8.legend(fontsize=8)

# Plot 9: 3D Bloch sphere trajectory (optimal)
ax9 = plt.subplot(3, 3, 9, projection='3d')
bloch_x_optimal = np.real([np.trace(sigma_x @ rho) for rho in rho_trajectory_optimal])
bloch_y_optimal = np.real([np.trace(sigma_y @ rho) for rho in rho_trajectory_optimal])
bloch_z_optimal = np.real([np.trace(sigma_z @ rho) for rho in rho_trajectory_optimal])

ax9.plot_surface(x_sphere, y_sphere, z_sphere, color='cyan', alpha=0.1)
ax9.plot(bloch_x_optimal, bloch_y_optimal, bloch_z_optimal, 'r-', linewidth=2, label='Trajectory')
ax9.scatter([bloch_x_optimal[0]], [bloch_y_optimal[0]], [bloch_z_optimal[0]], 
            color='green', s=100, label='Start')
ax9.scatter([bloch_x_optimal[-1]], [bloch_y_optimal[-1]], [bloch_z_optimal[-1]], 
            color='red', s=100, label='End')
ax9.set_xlabel('X', fontsize=10)
ax9.set_ylabel('Y', fontsize=10)
ax9.set_zlabel('Z', fontsize=10)
ax9.set_title('Optimal Control: Bloch Sphere', fontsize=14, fontweight='bold')
ax9.legend(fontsize=8)

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

print("\n" + "="*60)
print("SUMMARY")
print("="*60)
print(f"Final fidelity (Initial): {1-fid_err_initial:.4f}")
print(f"Final fidelity (Optimal): {1-fid_err_optimal:.4f}")
print(f"Dissipation reduction: {(diss_initial-diss_optimal)/diss_initial*100:.2f}%")
print(f"Control cost reduction: {(ctrl_cost_initial-ctrl_cost_optimal)/ctrl_cost_initial*100:.2f}%")
print(f"Total cost reduction: {(cost_initial-cost_optimal)/cost_initial*100:.2f}%")
print("="*60)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.integrate import quad
from scipy.linalg import expm

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

# Physical parameters
gamma = 1.0  # Coupling strength
omega_c = 10.0  # Cutoff frequency
T = 10.0  # Total evolution time
n_points = 200  # Number of time points for plotting

# Noise spectral density (1/f^alpha noise)
def spectral_density(omega, alpha=1.0):
    """
    Power spectral density of the noise
    S(omega) ~ 1/omega^alpha for omega > 0
    """
    if omega == 0:
        return 0
    return gamma**2 / (np.abs(omega)**alpha + 0.1)

# Filter function for different DD sequences
def filter_function(omega, pulse_times, T):
    """
    Calculate the filter function for a given pulse sequence
    pulse_times: array of pulse application times
    """
    N = len(pulse_times)
    
    # Add boundary points
    times = np.concatenate([[0], pulse_times, [T]])
    
    # Calculate filter function
    y = 0
    for i in range(len(times) - 1):
        t1 = times[i]
        t2 = times[i + 1]
        sign = (-1)**i
        
        if omega == 0:
            y += sign * (t2 - t1)
        else:
            y += sign * (np.sin(omega * t2) - np.sin(omega * t1)) / omega
    
    return np.abs(y)**2

# Calculate decoherence function
def calculate_decoherence(pulse_times, T, omega_max=50, n_omega=1000):
    """
    Calculate the decoherence by integrating over the noise spectrum
    """
    omegas = np.linspace(0.01, omega_max, n_omega)
    integrand = []
    
    for omega in omegas:
        F = filter_function(omega, pulse_times, T)
        S = spectral_density(omega)
        integrand.append(F * S)
    
    # Numerical integration
    chi = np.trapz(integrand, omegas)
    
    # Decoherence function (fidelity decay)
    return np.exp(-chi / 2)

# Generate different DD sequences
def free_evolution(T):
    """No pulses"""
    return np.array([])

def spin_echo(T):
    """Single pi pulse at T/2"""
    return np.array([T/2])

def cpmg_sequence(T, n_pulses):
    """CPMG sequence with n equally spaced pulses"""
    return np.linspace(T/(2*n_pulses), T - T/(2*n_pulses), n_pulses)

def udd_sequence(T, n_pulses):
    """Uhrig Dynamical Decoupling sequence"""
    k = np.arange(1, n_pulses + 1)
    return T * np.sin(np.pi * k / (2 * (n_pulses + 1)))**2

# Compare different sequences
print("Calculating decoherence for different DD sequences...")
print("=" * 60)

n_pulses_range = range(1, 11)
fidelities = {
    'Free Evolution': [],
    'CPMG': [],
    'UDD': []
}

for n in n_pulses_range:
    if n == 1:
        # Free evolution
        f_free = calculate_decoherence(free_evolution(T), T)
        fidelities['Free Evolution'].append(f_free)
        print(f"Free Evolution: Fidelity = {f_free:.6f}")
    else:
        fidelities['Free Evolution'].append(f_free)
    
    # CPMG
    pulses_cpmg = cpmg_sequence(T, n)
    f_cpmg = calculate_decoherence(pulses_cpmg, T)
    fidelities['CPMG'].append(f_cpmg)
    print(f"CPMG (n={n}): Fidelity = {f_cpmg:.6f}")
    
    # UDD
    pulses_udd = udd_sequence(T, n)
    f_udd = calculate_decoherence(pulses_udd, T)
    fidelities['UDD'].append(f_udd)
    print(f"UDD (n={n}): Fidelity = {f_udd:.6f}")

print("=" * 60)

# Visualization 1: Fidelity vs Number of Pulses
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.plot(n_pulses_range, fidelities['Free Evolution'], 'o-', 
         label='Free Evolution', linewidth=2, markersize=8)
plt.plot(n_pulses_range, fidelities['CPMG'], 's-', 
         label='CPMG', linewidth=2, markersize=8)
plt.plot(n_pulses_range, fidelities['UDD'], '^-', 
         label='UDD', linewidth=2, markersize=8)
plt.xlabel('Number of Pulses', fontsize=12)
plt.ylabel('Fidelity', fontsize=12)
plt.title('Decoherence Suppression Performance', fontsize=14, fontweight='bold')
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)

# Visualization 2: Pulse timing comparison
plt.subplot(1, 2, 2)
n_compare = 5
times_cpmg = cpmg_sequence(T, n_compare)
times_udd = udd_sequence(T, n_compare)

plt.scatter(times_cpmg, np.ones_like(times_cpmg), s=100, marker='s', 
           label='CPMG', alpha=0.7)
plt.scatter(times_udd, np.ones_like(times_udd) * 0.5, s=100, marker='^', 
           label='UDD', alpha=0.7)
plt.yticks([0.5, 1.0], ['UDD', 'CPMG'])
plt.xlabel('Time', fontsize=12)
plt.title(f'Pulse Timing Comparison (n={n_compare})', fontsize=14, fontweight='bold')
plt.xlim(0, T)
plt.ylim(0.2, 1.3)
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3, axis='x')

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

# Visualization 3: Filter Functions
print("\nCalculating filter functions...")
omega_range = np.linspace(0.1, 30, 300)
n_test = 5

filter_free = np.array([filter_function(w, free_evolution(T), T) for w in omega_range])
filter_cpmg = np.array([filter_function(w, cpmg_sequence(T, n_test), T) for w in omega_range])
filter_udd = np.array([filter_function(w, udd_sequence(T, n_test), T) for w in omega_range])

plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.semilogy(omega_range, filter_free, label='Free Evolution', linewidth=2)
plt.semilogy(omega_range, filter_cpmg, label=f'CPMG (n={n_test})', linewidth=2)
plt.semilogy(omega_range, filter_udd, label=f'UDD (n={n_test})', linewidth=2)
plt.xlabel('Frequency $\\omega$', fontsize=12)
plt.ylabel('Filter Function $F(\\omega)$', fontsize=12)
plt.title('Filter Functions (log scale)', fontsize=14, fontweight='bold')
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)

# Noise spectrum
plt.subplot(1, 2, 2)
noise_spectrum = np.array([spectral_density(w) for w in omega_range])
plt.semilogy(omega_range, noise_spectrum, 'r-', linewidth=2, label='Noise Spectrum')
plt.xlabel('Frequency $\\omega$', fontsize=12)
plt.ylabel('Spectral Density $S(\\omega)$', fontsize=12)
plt.title('Environmental Noise Spectrum', fontsize=14, fontweight='bold')
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)

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

# Visualization 4: 3D Surface - Fidelity vs (n_pulses, total_time)
print("\nGenerating 3D surface plot...")
n_pulse_3d = np.arange(1, 16)
time_3d = np.linspace(2, 20, 15)
N_grid, T_grid = np.meshgrid(n_pulse_3d, time_3d)

fidelity_cpmg_3d = np.zeros_like(N_grid, dtype=float)
fidelity_udd_3d = np.zeros_like(N_grid, dtype=float)

for i, t_val in enumerate(time_3d):
    for j, n_val in enumerate(n_pulse_3d):
        pulses_cpmg = cpmg_sequence(t_val, n_val)
        pulses_udd = udd_sequence(t_val, n_val)
        fidelity_cpmg_3d[i, j] = calculate_decoherence(pulses_cpmg, t_val)
        fidelity_udd_3d[i, j] = calculate_decoherence(pulses_udd, t_val)

fig = plt.figure(figsize=(16, 6))

# CPMG 3D surface
ax1 = fig.add_subplot(121, projection='3d')
surf1 = ax1.plot_surface(N_grid, T_grid, fidelity_cpmg_3d, cmap='viridis', 
                          alpha=0.9, edgecolor='none')
ax1.set_xlabel('Number of Pulses', fontsize=11, labelpad=10)
ax1.set_ylabel('Total Time', fontsize=11, labelpad=10)
ax1.set_zlabel('Fidelity', fontsize=11, labelpad=10)
ax1.set_title('CPMG Sequence Performance', fontsize=13, fontweight='bold', pad=20)
ax1.view_init(elev=25, azim=45)
fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=5)

# UDD 3D surface
ax2 = fig.add_subplot(122, projection='3d')
surf2 = ax2.plot_surface(N_grid, T_grid, fidelity_udd_3d, cmap='plasma', 
                          alpha=0.9, edgecolor='none')
ax2.set_xlabel('Number of Pulses', fontsize=11, labelpad=10)
ax2.set_ylabel('Total Time', fontsize=11, labelpad=10)
ax2.set_zlabel('Fidelity', fontsize=11, labelpad=10)
ax2.set_title('UDD Sequence Performance', fontsize=13, fontweight='bold', pad=20)
ax2.view_init(elev=25, azim=45)
fig.colorbar(surf2, ax=ax2, shrink=0.5, aspect=5)

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

# Summary statistics
print("\nSummary of Optimization Results:")
print("=" * 60)
optimal_n_cpmg = n_pulses_range[np.argmax(fidelities['CPMG'])]
optimal_n_udd = n_pulses_range[np.argmax(fidelities['UDD'])]
print(f"Optimal CPMG pulses: n = {optimal_n_cpmg}, Fidelity = {max(fidelities['CPMG']):.6f}")
print(f"Optimal UDD pulses: n = {optimal_n_udd}, Fidelity = {max(fidelities['UDD']):.6f}")
print(f"Free evolution fidelity: {fidelities['Free Evolution'][0]:.6f}")
print(f"\nImprovement factor (UDD vs Free): {max(fidelities['UDD']) / fidelities['Free Evolution'][0]:.3f}x")
print("=" * 60)

import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# Pauli matrices
sigma_x = np.array([[0, 1], [1, 0]], dtype=complex)
sigma_y = np.array([[0, -1j], [1j, 0]], dtype=complex)
sigma_z = np.array([[1, 0], [0, -1]], dtype=complex)
I = np.eye(2, dtype=complex)

class GRAPEOptimizer:
    def __init__(self, omega0, T, N, u_max):
        """
        Initialize GRAPE optimizer for qubit control
        
        Parameters:
        omega0: qubit frequency
        T: total evolution time
        N: number of time steps
        u_max: maximum control amplitude
        """
        self.omega0 = omega0
        self.T = T
        self.N = N
        self.dt = T / N
        self.u_max = u_max
        
        # Drift Hamiltonian (free evolution)
        self.H_drift = 0.5 * omega0 * sigma_z
        
        # Control Hamiltonian
        self.H_control = sigma_x
        
    def get_hamiltonian(self, u):
        """Construct total Hamiltonian for control amplitude u"""
        return self.H_drift + u * self.H_control
    
    def propagator(self, u):
        """Calculate propagator for single time step"""
        H = self.get_hamiltonian(u)
        return expm(-1j * H * self.dt)
    
    def forward_propagation(self, u_array, psi_init):
        """
        Forward propagate the state through all time steps
        Returns list of states at each time step
        """
        states = [psi_init]
        psi = psi_init.copy()
        
        for u in u_array:
            U = self.propagator(u)
            psi = U @ psi
            states.append(psi.copy())
        
        return states
    
    def calculate_fidelity(self, psi_final, psi_target):
        """Calculate fidelity between final and target states"""
        overlap = np.abs(np.vdot(psi_target, psi_final))**2
        return overlap
    
    def gradient(self, u_array, psi_init, psi_target):
        """
        Calculate gradient of fidelity with respect to control amplitudes
        Using adjoint method for efficient computation
        """
        # Forward propagation
        forward_states = self.forward_propagation(u_array, psi_init)
        
        # Backward propagation (adjoint)
        psi_target_conj = psi_target.conj()
        lambda_states = [psi_target_conj]
        
        for k in range(self.N-1, -1, -1):
            U = self.propagator(u_array[k])
            lambda_k = U.T.conj() @ lambda_states[0]
            lambda_states.insert(0, lambda_k)
        
        # Calculate gradient
        grad = np.zeros(self.N)
        for k in range(self.N):
            psi_k = forward_states[k]
            lambda_k = lambda_states[k+1]
            
            # Derivative of propagator with respect to u
            dU_du = -1j * self.dt * self.H_control @ self.propagator(u_array[k])
            
            # Gradient component
            grad[k] = 2 * np.real(np.vdot(lambda_k, dU_du @ psi_k))
        
        return grad
    
    def optimize(self, psi_init, psi_target, max_iter=200, learning_rate=0.5, tolerance=1e-6):
        """
        Optimize control pulse using GRAPE algorithm
        """
        # Initialize control pulse (random small values)
        u_array = np.random.randn(self.N) * 0.1
        
        fidelities = []
        
        for iteration in range(max_iter):
            # Calculate current fidelity
            final_states = self.forward_propagation(u_array, psi_init)
            psi_final = final_states[-1]
            fidelity = self.calculate_fidelity(psi_final, psi_target)
            fidelities.append(fidelity)
            
            if iteration % 20 == 0:
                print(f"Iteration {iteration}: Fidelity = {fidelity:.6f}")
            
            # Check convergence
            if fidelity > 1 - tolerance:
                print(f"Converged at iteration {iteration}")
                break
            
            # Calculate gradient
            grad = self.gradient(u_array, psi_init, psi_target)
            
            # Update control with gradient ascent
            u_array = u_array + learning_rate * grad
            
            # Apply amplitude constraints
            u_array = np.clip(u_array, -self.u_max, self.u_max)
        
        return u_array, fidelities

# Set up the problem
omega0 = 2 * np.pi * 1.0  # 1 GHz qubit frequency
T = 10.0  # Total time (arbitrary units)
N = 100  # Number of time steps
u_max = 2.0  # Maximum control amplitude

# Initial and target states
psi_init = np.array([1, 0], dtype=complex)  # Ground state |0>
psi_target = np.array([0, 1], dtype=complex)  # Excited state |1>

# Create optimizer and run GRAPE
print("Starting GRAPE optimization...")
print(f"Qubit frequency: {omega0/(2*np.pi):.2f} GHz")
print(f"Total time: {T}")
print(f"Time steps: {N}")
print(f"Max control amplitude: {u_max}")
print()

optimizer = GRAPEOptimizer(omega0, T, N, u_max)
u_optimal, fidelities = optimizer.optimize(psi_init, psi_target, max_iter=200, learning_rate=0.3)

print()
print(f"Final fidelity: {fidelities[-1]:.8f}")

# Verify the result
final_states = optimizer.forward_propagation(u_optimal, psi_init)
psi_final = final_states[-1]
print(f"Final state: {psi_final}")
print(f"Target state: {psi_target}")

# Calculate Bloch sphere trajectory
def state_to_bloch(psi):
    """Convert quantum state to Bloch sphere coordinates"""
    x = 2 * np.real(psi[0].conj() * psi[1])
    y = 2 * np.imag(psi[0].conj() * psi[1])
    z = np.abs(psi[0])**2 - np.abs(psi[1])**2
    return x, y, z

bloch_trajectory = [state_to_bloch(state) for state in final_states]
bloch_x = [b[0] for b in bloch_trajectory]
bloch_y = [b[1] for b in bloch_trajectory]
bloch_z = [b[2] for b in bloch_trajectory]

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

# 1. Control pulse shape
ax1 = plt.subplot(3, 3, 1)
time_points = np.linspace(0, T, N)
ax1.plot(time_points, u_optimal, 'b-', linewidth=2)
ax1.axhline(y=0, color='k', linestyle='--', alpha=0.3)
ax1.set_xlabel('Time', fontsize=12)
ax1.set_ylabel('Control Amplitude u(t)', fontsize=12)
ax1.set_title('Optimized Control Pulse', fontsize=14, fontweight='bold')
ax1.grid(True, alpha=0.3)

# 2. Fidelity evolution
ax2 = plt.subplot(3, 3, 2)
ax2.plot(fidelities, 'r-', linewidth=2)
ax2.set_xlabel('Iteration', fontsize=12)
ax2.set_ylabel('Fidelity', fontsize=12)
ax2.set_title('Fidelity vs Iteration', fontsize=14, fontweight='bold')
ax2.grid(True, alpha=0.3)
ax2.set_ylim([0, 1.05])

# 3. Population dynamics
ax3 = plt.subplot(3, 3, 3)
pop_0 = [np.abs(state[0])**2 for state in final_states]
pop_1 = [np.abs(state[1])**2 for state in final_states]
time_evolution = np.linspace(0, T, N+1)
ax3.plot(time_evolution, pop_0, 'b-', linewidth=2, label='|0⟩')
ax3.plot(time_evolution, pop_1, 'r-', linewidth=2, label='|1⟩')
ax3.set_xlabel('Time', fontsize=12)
ax3.set_ylabel('Population', fontsize=12)
ax3.set_title('State Population Dynamics', fontsize=14, fontweight='bold')
ax3.legend(fontsize=11)
ax3.grid(True, alpha=0.3)

# 4. Bloch sphere trajectory (3D)
ax4 = plt.subplot(3, 3, 4, projection='3d')

# Draw Bloch sphere
u_sphere = np.linspace(0, 2 * np.pi, 50)
v_sphere = np.linspace(0, np.pi, 50)
x_sphere = np.outer(np.cos(u_sphere), np.sin(v_sphere))
y_sphere = np.outer(np.sin(u_sphere), np.sin(v_sphere))
z_sphere = np.outer(np.ones(np.size(u_sphere)), np.cos(v_sphere))
ax4.plot_surface(x_sphere, y_sphere, z_sphere, alpha=0.1, color='cyan')

# Plot trajectory
ax4.plot(bloch_x, bloch_y, bloch_z, 'b-', linewidth=2, label='Trajectory')
ax4.scatter([bloch_x[0]], [bloch_y[0]], [bloch_z[0]], color='green', s=100, label='Start')
ax4.scatter([bloch_x[-1]], [bloch_y[-1]], [bloch_z[-1]], color='red', s=100, label='End')

ax4.set_xlabel('X', fontsize=10)
ax4.set_ylabel('Y', fontsize=10)
ax4.set_zlabel('Z', fontsize=10)
ax4.set_title('Bloch Sphere Trajectory', fontsize=14, fontweight='bold')
ax4.legend(fontsize=9)

# 5. Bloch X-Y projection
ax5 = plt.subplot(3, 3, 5)
ax5.plot(bloch_x, bloch_y, 'b-', linewidth=2)
ax5.scatter([bloch_x[0]], [bloch_y[0]], color='green', s=100, zorder=5, label='Start')
ax5.scatter([bloch_x[-1]], [bloch_y[-1]], color='red', s=100, zorder=5, label='End')
circle = plt.Circle((0, 0), 1, fill=False, color='gray', linestyle='--', alpha=0.5)
ax5.add_patch(circle)
ax5.set_xlabel('X', fontsize=12)
ax5.set_ylabel('Y', fontsize=12)
ax5.set_title('Bloch Sphere (X-Y Projection)', fontsize=14, fontweight='bold')
ax5.set_aspect('equal')
ax5.grid(True, alpha=0.3)
ax5.legend(fontsize=10)

# 6. Bloch X-Z projection
ax6 = plt.subplot(3, 3, 6)
ax6.plot(bloch_x, bloch_z, 'b-', linewidth=2)
ax6.scatter([bloch_x[0]], [bloch_z[0]], color='green', s=100, zorder=5, label='Start')
ax6.scatter([bloch_x[-1]], [bloch_z[-1]], color='red', s=100, zorder=5, label='End')
circle = plt.Circle((0, 0), 1, fill=False, color='gray', linestyle='--', alpha=0.5)
ax6.add_patch(circle)
ax6.set_xlabel('X', fontsize=12)
ax6.set_ylabel('Z', fontsize=12)
ax6.set_title('Bloch Sphere (X-Z Projection)', fontsize=14, fontweight='bold')
ax6.set_aspect('equal')
ax6.grid(True, alpha=0.3)
ax6.legend(fontsize=10)

# 7. Control pulse frequency spectrum
ax7 = plt.subplot(3, 3, 7)
fft_u = np.fft.fft(u_optimal)
freqs = np.fft.fftfreq(N, d=optimizer.dt)
power_spectrum = np.abs(fft_u)**2
ax7.plot(freqs[:N//2], power_spectrum[:N//2], 'purple', linewidth=2)
ax7.set_xlabel('Frequency', fontsize=12)
ax7.set_ylabel('Power', fontsize=12)
ax7.set_title('Control Pulse Frequency Spectrum', fontsize=14, fontweight='bold')
ax7.grid(True, alpha=0.3)

# 8. Phase evolution
ax8 = plt.subplot(3, 3, 8)
phases = [np.angle(state[1]) if np.abs(state[1]) > 1e-10 else 0 for state in final_states]
ax8.plot(time_evolution, phases, 'orange', linewidth=2)
ax8.set_xlabel('Time', fontsize=12)
ax8.set_ylabel('Phase (rad)', fontsize=12)
ax8.set_title('Excited State Phase Evolution', fontsize=14, fontweight='bold')
ax8.grid(True, alpha=0.3)

# 9. Control energy
ax9 = plt.subplot(3, 3, 9)
energy = np.cumsum(u_optimal**2) * optimizer.dt
ax9.plot(time_points, energy, 'brown', linewidth=2)
ax9.set_xlabel('Time', fontsize=12)
ax9.set_ylabel('Cumulative Energy', fontsize=12)
ax9.set_title('Control Energy Consumption', fontsize=14, fontweight='bold')
ax9.grid(True, alpha=0.3)

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

# Additional 3D visualization: Control pulse in time-frequency domain
fig2 = plt.figure(figsize=(14, 6))

# 3D plot: Time-Frequency-Amplitude
ax_3d = fig2.add_subplot(121, projection='3d')

# Create time-frequency representation using sliding window FFT
window_size = 20
hop_size = 5
n_windows = (N - window_size) // hop_size + 1

time_centers = []
freq_arrays = []
magnitude_arrays = []

for i in range(n_windows):
    start = i * hop_size
    end = start + window_size
    window = u_optimal[start:end]
    
    fft_window = np.fft.fft(window)
    freqs_window = np.fft.fftfreq(window_size, d=optimizer.dt)
    
    time_centers.append(time_points[start + window_size//2])
    freq_arrays.append(freqs_window[:window_size//2])
    magnitude_arrays.append(np.abs(fft_window[:window_size//2]))

# Create meshgrid for 3D surface
T_mesh = np.array(time_centers)
F_mesh = np.array(freq_arrays[0])
T_grid, F_grid = np.meshgrid(T_mesh, F_mesh)
Z_grid = np.array(magnitude_arrays).T

surf = ax_3d.plot_surface(T_grid, F_grid, Z_grid, cmap=cm.viridis, alpha=0.8)
ax_3d.set_xlabel('Time', fontsize=11)
ax_3d.set_ylabel('Frequency', fontsize=11)
ax_3d.set_zlabel('Magnitude', fontsize=11)
ax_3d.set_title('Time-Frequency Analysis (3D)', fontsize=13, fontweight='bold')
fig2.colorbar(surf, ax=ax_3d, shrink=0.5)

# 2D heatmap version
ax_heat = fig2.add_subplot(122)
im = ax_heat.contourf(T_grid, F_grid, Z_grid, levels=20, cmap=cm.viridis)
ax_heat.set_xlabel('Time', fontsize=12)
ax_heat.set_ylabel('Frequency', fontsize=12)
ax_heat.set_title('Time-Frequency Heatmap', fontsize=13, fontweight='bold')
fig2.colorbar(im, ax=ax_heat)

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

print("\n" + "="*60)
print("GRAPE Optimization Complete!")
print("="*60)
print(f"Final fidelity achieved: {fidelities[-1]:.8f}")
print(f"Total control energy: {np.sum(u_optimal**2) * optimizer.dt:.4f}")
print(f"Peak control amplitude: {np.max(np.abs(u_optimal)):.4f}")
print("Figures saved: 'grape_optimization_results.png' and 'grape_time_frequency_3d.png'")

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

# Physical constants and parameters
hbar = 1.0
omega_0 = 2.0 * np.pi
T = 1.0
N_steps = 50
dt = T / N_steps
t_array = np.linspace(0, T, N_steps + 1)

# Pauli matrices
sigma_x = np.array([[0, 1], [1, 0]], dtype=complex)
sigma_z = np.array([[1, 0], [0, -1]], dtype=complex)

# Initial and target states
psi_initial = np.array([1, 0], dtype=complex)
psi_target = np.array([1, 1], dtype=complex) / np.sqrt(2)

def evolve_state_fast(u_values):
    """Fast state evolution using matrix exponential"""
    psi = psi_initial.copy()
    psi_evolution = [psi.copy()]
    
    for i in range(N_steps):
        H = (omega_0 / 2) * sigma_z + u_values[i] * sigma_x
        U = expm(-1j * H * dt / hbar)
        psi = np.dot(U, psi)
        psi = psi / np.linalg.norm(psi)
        psi_evolution.append(psi.copy())
    
    return psi, np.array(psi_evolution).T

def fidelity(psi_final, psi_target):
    """Calculate fidelity"""
    return np.abs(np.vdot(psi_target, psi_final))**2

def objective_function(u_values):
    """Objective function to minimize"""
    psi_final, _ = evolve_state_fast(u_values)
    F = fidelity(psi_final, psi_target)
    regularization = 0.001 * np.sum(u_values**2) * dt
    return -(F - regularization)

# Improved initial guess using analytical insight
u_initial = np.sin(np.pi * t_array[:-1] / T) * 3.0

print("Starting optimization...")
print(f"Initial fidelity: {-objective_function(u_initial):.6f}")

# Optimization with reduced iterations
result = minimize(
    objective_function,
    u_initial,
    method='Powell',
    options={'maxiter': 500, 'disp': True}
)

u_optimal = result.x
psi_final_optimal, psi_evolution = evolve_state_fast(u_optimal)
final_fidelity = fidelity(psi_final_optimal, psi_target)

print(f"\nOptimization completed!")
print(f"Final fidelity: {final_fidelity:.6f}")
print(f"Error: {1 - final_fidelity:.6e}")

# Bloch sphere coordinate conversion
def state_to_bloch(psi):
    x = 2 * np.real(psi[0] * np.conj(psi[1]))
    y = 2 * np.imag(psi[0] * np.conj(psi[1]))
    z = np.abs(psi[0])**2 - np.abs(psi[1])**2
    return x, y, z

bloch_coords = np.array([state_to_bloch(psi_evolution[:, i]) for i in range(N_steps + 1)])

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

# Plot 1: Optimal control pulse
ax1 = plt.subplot(3, 3, 1)
plt.plot(t_array[:-1], u_optimal, 'b-', linewidth=2)
plt.xlabel('Time', fontsize=12)
plt.ylabel('Control amplitude u(t)', fontsize=12)
plt.title('Optimal Control Pulse', fontsize=14, fontweight='bold')
plt.grid(True, alpha=0.3)

# Plot 2: Fidelity evolution
fidelity_evolution = [fidelity(psi_evolution[:, i], psi_target) for i in range(N_steps + 1)]

ax2 = plt.subplot(3, 3, 2)
plt.plot(t_array, fidelity_evolution, 'r-', linewidth=2)
plt.xlabel('Time', fontsize=12)
plt.ylabel('Fidelity', fontsize=12)
plt.title('Fidelity Evolution', fontsize=14, fontweight='bold')
plt.grid(True, alpha=0.3)
plt.ylim([0, 1.05])

# Plot 3: State populations
populations_0 = np.abs(psi_evolution[0, :])**2
populations_1 = np.abs(psi_evolution[1, :])**2

ax3 = plt.subplot(3, 3, 3)
plt.plot(t_array, populations_0, 'b-', linewidth=2, label='|0⟩')
plt.plot(t_array, populations_1, 'r-', linewidth=2, label='|1⟩')
plt.xlabel('Time', fontsize=12)
plt.ylabel('Population', fontsize=12)
plt.title('State Populations', fontsize=14, fontweight='bold')
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)
plt.ylim([0, 1.05])

# Plot 4: Bloch sphere trajectory (3D)
ax4 = fig.add_subplot(3, 3, 4, projection='3d')

u_sphere = np.linspace(0, 2 * np.pi, 30)
v_sphere = np.linspace(0, np.pi, 30)
x_sphere = np.outer(np.cos(u_sphere), np.sin(v_sphere))
y_sphere = np.outer(np.sin(u_sphere), np.sin(v_sphere))
z_sphere = np.outer(np.ones(np.size(u_sphere)), np.cos(v_sphere))
ax4.plot_surface(x_sphere, y_sphere, z_sphere, alpha=0.1, color='cyan')

ax4.plot(bloch_coords[:, 0], bloch_coords[:, 1], bloch_coords[:, 2], 
         'b-', linewidth=2, label='Trajectory')
ax4.scatter([bloch_coords[0, 0]], [bloch_coords[0, 1]], [bloch_coords[0, 2]], 
            c='green', s=100, marker='o', label='Initial')
ax4.scatter([bloch_coords[-1, 0]], [bloch_coords[-1, 1]], [bloch_coords[-1, 2]], 
            c='red', s=100, marker='*', label='Final')

x_target, y_target, z_target = state_to_bloch(psi_target)
ax4.scatter([x_target], [y_target], [z_target], 
            c='orange', s=150, marker='X', label='Target', edgecolors='black', linewidths=2)

ax4.set_xlabel('X', fontsize=11)
ax4.set_ylabel('Y', fontsize=11)
ax4.set_zlabel('Z', fontsize=11)
ax4.set_title('Bloch Sphere Trajectory', fontsize=14, fontweight='bold')
ax4.legend(fontsize=9)
ax4.set_box_aspect([1,1,1])

# Plot 5: State amplitude components
ax5 = plt.subplot(3, 3, 5)
plt.plot(t_array, psi_evolution[0, :].real, 'b-', linewidth=2, label='Re(c₀)')
plt.plot(t_array, psi_evolution[0, :].imag, 'b--', linewidth=2, label='Im(c₀)')
plt.plot(t_array, psi_evolution[1, :].real, 'r-', linewidth=2, label='Re(c₁)')
plt.plot(t_array, psi_evolution[1, :].imag, 'r--', linewidth=2, label='Im(c₁)')
plt.xlabel('Time', fontsize=12)
plt.ylabel('Amplitude', fontsize=12)
plt.title('State Amplitude Components', fontsize=14, fontweight='bold')
plt.legend(fontsize=9, ncol=2)
plt.grid(True, alpha=0.3)

# Plot 6: Phase evolution
phases_0 = np.angle(psi_evolution[0, :])
phases_1 = np.angle(psi_evolution[1, :])

ax6 = plt.subplot(3, 3, 6)
plt.plot(t_array, phases_0, 'b-', linewidth=2, label='Phase(c₀)')
plt.plot(t_array, phases_1, 'r-', linewidth=2, label='Phase(c₁)')
plt.xlabel('Time', fontsize=12)
plt.ylabel('Phase (radians)', fontsize=12)
plt.title('Phase Evolution', fontsize=14, fontweight='bold')
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)

# Plot 7: Control pulse power spectrum
from scipy.fft import fft, fftfreq
u_padded = np.pad(u_optimal, (0, 512 - len(u_optimal)), mode='constant')
freqs = fftfreq(len(u_padded), dt)
spectrum = np.abs(fft(u_padded))

ax7 = plt.subplot(3, 3, 7)
plt.plot(freqs[:len(freqs)//2], spectrum[:len(spectrum)//2], 'purple', linewidth=2)
plt.xlabel('Frequency', fontsize=12)
plt.ylabel('Magnitude', fontsize=12)
plt.title('Control Pulse Frequency Spectrum', fontsize=14, fontweight='bold')
plt.grid(True, alpha=0.3)

# Plot 8: Energy evolution
energy = []
for i in range(N_steps):
    psi_t = psi_evolution[:, i]
    H = (omega_0 / 2) * sigma_z + u_optimal[i] * sigma_x
    E = np.real(np.vdot(psi_t, np.dot(H, psi_t)))
    energy.append(E)

ax8 = plt.subplot(3, 3, 8)
plt.plot(t_array[:-1], energy, 'g-', linewidth=2)
plt.xlabel('Time', fontsize=12)
plt.ylabel('Energy ⟨H⟩', fontsize=12)
plt.title('System Energy Evolution', fontsize=14, fontweight='bold')
plt.grid(True, alpha=0.3)

# Plot 9: 3D trajectory in state space
ax9 = fig.add_subplot(3, 3, 9, projection='3d')
ax9.plot(np.abs(psi_evolution[0, :]), np.abs(psi_evolution[1, :]), t_array, 
         'b-', linewidth=2)
ax9.scatter([np.abs(psi_evolution[0, 0])], [np.abs(psi_evolution[1, 0])], [0], 
            c='green', s=100, marker='o', label='Initial')
ax9.scatter([np.abs(psi_evolution[0, -1])], [np.abs(psi_evolution[1, -1])], [T], 
            c='red', s=100, marker='*', label='Final')
ax9.set_xlabel('|c₀|', fontsize=11)
ax9.set_ylabel('|c₁|', fontsize=11)
ax9.set_zlabel('Time', fontsize=11)
ax9.set_title('State Evolution in Amplitude Space', fontsize=14, fontweight='bold')
ax9.legend(fontsize=9)

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

# Print final results
print("\n" + "="*60)
print("OPTIMIZATION RESULTS")
print("="*60)
print(f"Initial state: |ψ₀⟩ = |0⟩")
print(f"Target state: |ψₜ⟩ = (|0⟩ + |1⟩)/√2")
print(f"Evolution time: T = {T}")
print(f"Number of control steps: {N_steps}")
print(f"Final fidelity: {final_fidelity:.8f}")
print(f"Fidelity error: {1 - final_fidelity:.2e}")
print(f"\nFinal state coefficients:")
print(f"  c₀ = {psi_final_optimal[0]:.6f}")
print(f"  c₁ = {psi_final_optimal[1]:.6f}")
print(f"\nTarget state coefficients:")
print(f"  c₀ = {psi_target[0]:.6f}")
print(f"  c₁ = {psi_target[1]:.6f}")
print(f"\nControl pulse statistics:")
print(f"  Maximum amplitude: {np.max(np.abs(u_optimal)):.4f}")
print(f"  Mean amplitude: {np.mean(np.abs(u_optimal)):.4f}")
print(f"  Control energy: {np.sum(u_optimal**2) * dt:.4f}")
print("="*60)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.optimize import minimize
import time

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

# ============================================
# VQE Implementation
# ============================================

class VQE:
    def __init__(self, hamiltonian_coeffs):
        """
        Initialize VQE solver
        hamiltonian_coeffs: dictionary with Pauli string keys and coefficient values
        Example: {'Z0': 0.5, 'Z1': 0.8, 'X0X1': 0.3}
        """
        self.hamiltonian = hamiltonian_coeffs
        self.optimization_history = []
        
    def apply_rotation(self, state, angle, pauli_type, qubit_idx):
        """Apply single-qubit rotation gate"""
        if pauli_type == 'X':
            rotation = np.array([[np.cos(angle/2), -1j*np.sin(angle/2)],
                                [-1j*np.sin(angle/2), np.cos(angle/2)]])
        elif pauli_type == 'Y':
            rotation = np.array([[np.cos(angle/2), -np.sin(angle/2)],
                                [np.sin(angle/2), np.cos(angle/2)]])
        elif pauli_type == 'Z':
            rotation = np.array([[np.exp(-1j*angle/2), 0],
                                [0, np.exp(1j*angle/2)]])
        
        # Apply rotation to specific qubit
        n_qubits = int(np.log2(len(state)))
        identity = np.eye(2)
        
        if qubit_idx == 0:
            gate = rotation
            for i in range(1, n_qubits):
                gate = np.kron(gate, identity)
        else:
            gate = identity
            for i in range(1, n_qubits):
                if i == qubit_idx:
                    gate = np.kron(gate, rotation)
                else:
                    gate = np.kron(gate, identity)
        
        return gate @ state
    
    def create_ansatz(self, params, n_qubits=2, depth=2):
        """Create parameterized quantum circuit (ansatz)"""
        # Initialize state |00...0>
        state = np.zeros(2**n_qubits, dtype=complex)
        state[0] = 1.0
        
        param_idx = 0
        for d in range(depth):
            # Rotation layer
            for q in range(n_qubits):
                state = self.apply_rotation(state, params[param_idx], 'Y', q)
                param_idx += 1
                state = self.apply_rotation(state, params[param_idx], 'Z', q)
                param_idx += 1
            
            # Entanglement layer (CNOT gates simulation)
            if d < depth - 1:
                for q in range(n_qubits - 1):
                    # CNOT between qubit q and q+1
                    cnot = np.array([[1, 0, 0, 0],
                                    [0, 1, 0, 0],
                                    [0, 0, 0, 1],
                                    [0, 0, 1, 0]])
                    if n_qubits > 2:
                        # Extend CNOT to full Hilbert space
                        full_cnot = np.eye(2**n_qubits)
                        # This is simplified; for full implementation use tensor products
                    state_reshaped = state.reshape([2] * n_qubits)
                    # Apply CNOT (simplified for 2 qubits)
                    if n_qubits == 2:
                        state = cnot @ state
        
        return state
    
    def measure_expectation(self, state, operator_string):
        """Measure expectation value of Pauli operator"""
        n_qubits = int(np.log2(len(state)))
        
        # Build Pauli operator matrix
        pauli_matrices = {
            'I': np.array([[1, 0], [0, 1]]),
            'X': np.array([[0, 1], [1, 0]]),
            'Y': np.array([[0, -1j], [1j, 0]]),
            'Z': np.array([[1, 0], [0, -1]])
        }
        
        # Parse operator string
        operator = None
        for i in range(n_qubits):
            qubit_op = 'I'
            for op_char in operator_string:
                if str(i) in operator_string:
                    if 'X' + str(i) in operator_string:
                        qubit_op = 'X'
                    elif 'Y' + str(i) in operator_string:
                        qubit_op = 'Y'
                    elif 'Z' + str(i) in operator_string:
                        qubit_op = 'Z'
            
            if operator is None:
                operator = pauli_matrices[qubit_op]
            else:
                operator = np.kron(operator, pauli_matrices[qubit_op])
        
        expectation = np.real(np.conj(state) @ operator @ state)
        return expectation
    
    def cost_function(self, params):
        """Calculate energy expectation value"""
        state = self.create_ansatz(params)
        
        energy = 0.0
        for pauli_string, coeff in self.hamiltonian.items():
            energy += coeff * self.measure_expectation(state, pauli_string)
        
        self.optimization_history.append(energy)
        return energy
    
    def optimize(self, initial_params=None, method='COBYLA', maxiter=100):
        """Run classical optimization"""
        if initial_params is None:
            initial_params = np.random.uniform(0, 2*np.pi, 8)
        
        self.optimization_history = []
        
        result = minimize(
            self.cost_function,
            initial_params,
            method=method,
            options={'maxiter': maxiter}
        )
        
        return result

# ============================================
# QAOA Implementation
# ============================================

class QAOA:
    def __init__(self, graph_edges):
        """
        Initialize QAOA solver for Max-Cut problem
        graph_edges: list of tuples representing edges [(0,1), (1,2), ...]
        """
        self.edges = graph_edges
        self.n_qubits = max([max(edge) for edge in graph_edges]) + 1
        self.optimization_history = []
        
    def apply_mixer(self, state, beta):
        """Apply mixer Hamiltonian (X rotations)"""
        for q in range(self.n_qubits):
            # Apply Rx(2*beta) to each qubit
            angle = 2 * beta
            cos_half = np.cos(angle / 2)
            sin_half = np.sin(angle / 2)
            
            # Build rotation matrix for full state
            new_state = np.zeros_like(state)
            for idx in range(len(state)):
                # Convert index to binary representation
                binary = format(idx, f'0{self.n_qubits}b')
                
                # Flip qubit q
                binary_list = list(binary)
                binary_list[q] = '1' if binary_list[q] == '0' else '0'
                flipped_idx = int(''.join(binary_list), 2)
                
                new_state[idx] += cos_half * state[idx]
                new_state[flipped_idx] += -1j * sin_half * state[idx]
            
            state = new_state
        
        return state
    
    def apply_cost(self, state, gamma):
        """Apply cost Hamiltonian (ZZ interactions)"""
        for edge in self.edges:
            i, j = edge
            # Apply exp(-i * gamma * Z_i Z_j)
            new_state = np.zeros_like(state)
            
            for idx in range(len(state)):
                binary = format(idx, f'0{self.n_qubits}b')
                
                # Check bits at positions i and j
                bit_i = int(binary[i])
                bit_j = int(binary[j])
                
                # ZZ eigenvalue: +1 if bits are same, -1 if different
                zz_eigenvalue = 1 if bit_i == bit_j else -1
                
                phase = np.exp(-1j * gamma * zz_eigenvalue)
                new_state[idx] = phase * state[idx]
            
            state = new_state
        
        return state
    
    def create_qaoa_circuit(self, params, p=1):
        """Create QAOA circuit with p layers"""
        # Initialize superposition state
        state = np.ones(2**self.n_qubits, dtype=complex) / np.sqrt(2**self.n_qubits)
        
        # Apply p layers
        for layer in range(p):
            gamma = params[2*layer]
            beta = params[2*layer + 1]
            
            state = self.apply_cost(state, gamma)
            state = self.apply_mixer(state, beta)
        
        return state
    
    def compute_cost(self, state):
        """Compute Max-Cut objective function"""
        cost = 0.0
        
        for edge in self.edges:
            i, j = edge
            
            # Compute expectation of (1 - Z_i Z_j) / 2
            expectation = 0.0
            for idx in range(len(state)):
                binary = format(idx, f'0{self.n_qubits}b')
                bit_i = int(binary[i])
                bit_j = int(binary[j])
                
                zz_eigenvalue = 1 if bit_i == bit_j else -1
                expectation += np.abs(state[idx])**2 * zz_eigenvalue
            
            cost += (1 - expectation) / 2
        
        return cost
    
    def cost_function(self, params):
        """QAOA cost function"""
        state = self.create_qaoa_circuit(params)
        cost = self.compute_cost(state)
        self.optimization_history.append(-cost)  # Negative for minimization
        return -cost
    
    def optimize(self, initial_params=None, p=1, method='COBYLA', maxiter=100):
        """Run classical optimization"""
        if initial_params is None:
            initial_params = np.random.uniform(0, np.pi, 2*p)
        
        self.optimization_history = []
        
        result = minimize(
            self.cost_function,
            initial_params,
            method=method,
            options={'maxiter': maxiter}
        )
        
        return result

# ============================================
# Visualization Functions
# ============================================

def plot_optimization_landscape_3d(cost_func, param_ranges, title):
    """Plot 3D optimization landscape"""
    fig = plt.figure(figsize=(12, 5))
    
    # 3D surface plot
    ax1 = fig.add_subplot(121, projection='3d')
    
    x_range = np.linspace(param_ranges[0][0], param_ranges[0][1], 30)
    y_range = np.linspace(param_ranges[1][0], param_ranges[1][1], 30)
    X, Y = np.meshgrid(x_range, y_range)
    
    Z = np.zeros_like(X)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Z[i, j] = cost_func([X[i, j], Y[i, j]])
    
    surf = ax1.plot_surface(X, Y, Z, cmap='viridis', alpha=0.8)
    ax1.set_xlabel('Parameter 1', fontsize=10)
    ax1.set_ylabel('Parameter 2', fontsize=10)
    ax1.set_zlabel('Cost', fontsize=10)
    ax1.set_title(f'{title}\n3D Landscape', fontsize=12)
    fig.colorbar(surf, ax=ax1, shrink=0.5)
    
    # 2D contour plot
    ax2 = fig.add_subplot(122)
    contour = ax2.contour(X, Y, Z, levels=20, cmap='viridis')
    ax2.clabel(contour, inline=True, fontsize=8)
    ax2.set_xlabel('Parameter 1', fontsize=10)
    ax2.set_ylabel('Parameter 2', fontsize=10)
    ax2.set_title(f'{title}\nContour Plot', fontsize=12)
    plt.colorbar(contour, ax=ax2)
    
    plt.tight_layout()
    plt.savefig('/tmp/landscape.png', dpi=150, bbox_inches='tight')
    plt.show()

def plot_convergence(histories, labels, title):
    """Plot optimization convergence"""
    plt.figure(figsize=(10, 6))
    
    for history, label in zip(histories, labels):
        plt.plot(history, marker='o', markersize=4, label=label, linewidth=2)
    
    plt.xlabel('Iteration', fontsize=12)
    plt.ylabel('Cost Function Value', fontsize=12)
    plt.title(title, fontsize=14, fontweight='bold')
    plt.legend(fontsize=10)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('/tmp/convergence.png', dpi=150, bbox_inches='tight')
    plt.show()

def plot_parameter_evolution(histories, title):
    """Plot parameter evolution during optimization"""
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle(title, fontsize=14, fontweight='bold')
    
    for idx, (history, ax) in enumerate(zip(histories, axes.flat)):
        iterations = range(len(history))
        ax.plot(iterations, history, marker='o', markersize=4, linewidth=2, color=f'C{idx}')
        ax.set_xlabel('Iteration', fontsize=10)
        ax.set_ylabel(f'Parameter {idx+1}', fontsize=10)
        ax.set_title(f'Parameter {idx+1} Evolution', fontsize=11)
        ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('/tmp/param_evolution.png', dpi=150, bbox_inches='tight')
    plt.show()

# ============================================
# Main Execution
# ============================================

print("="*70)
print("VARIATIONAL QUANTUM ALGORITHMS: VQE AND QAOA")
print("="*70)

# ============================================
# Part 1: VQE Example
# ============================================

print("\n" + "="*70)
print("PART 1: Variational Quantum Eigensolver (VQE)")
print("="*70)

# Define Hamiltonian: H = 0.5*Z0 + 0.8*Z1 + 0.3*X0*X1
hamiltonian = {
    'Z0': 0.5,
    'Z1': 0.8,
    'X0X1': 0.3
}

print("\nHamiltonian:")
print("H = 0.5*Z0 + 0.8*Z1 + 0.3*X0*X1")

# Create VQE instance
vqe = VQE(hamiltonian)

# Run optimization
print("\nRunning VQE optimization...")
start_time = time.time()
result_vqe = vqe.optimize(maxiter=150)
vqe_time = time.time() - start_time

print(f"\nOptimization completed in {vqe_time:.2f} seconds")
print(f"Ground state energy: {result_vqe.fun:.6f}")
print(f"Optimal parameters: {result_vqe.x}")
print(f"Number of iterations: {len(vqe.optimization_history)}")

# ============================================
# Part 2: QAOA Example
# ============================================

print("\n" + "="*70)
print("PART 2: Quantum Approximate Optimization Algorithm (QAOA)")
print("="*70)

# Define graph for Max-Cut problem (triangle graph)
edges = [(0, 1), (1, 2), (2, 0)]

print("\nMax-Cut Problem on Triangle Graph")
print(f"Edges: {edges}")

# Create QAOA instance
qaoa = QAOA(edges)

# Run optimization
print("\nRunning QAOA optimization (p=2)...")
start_time = time.time()
result_qaoa = qaoa.optimize(p=2, maxiter=150)
qaoa_time = time.time() - start_time

print(f"\nOptimization completed in {qaoa_time:.2f} seconds")
print(f"Maximum cut value: {-result_qaoa.fun:.6f}")
print(f"Optimal parameters: {result_qaoa.x}")
print(f"Number of iterations: {len(qaoa.optimization_history)}")

# ============================================
# Visualization
# ============================================

print("\n" + "="*70)
print("GENERATING VISUALIZATIONS")
print("="*70)

# Plot convergence comparison
plot_convergence(
    [vqe.optimization_history, qaoa.optimization_history],
    ['VQE', 'QAOA'],
    'Optimization Convergence: VQE vs QAOA'
)

# Create simplified 2D cost landscape for VQE
print("\nGenerating VQE cost landscape...")
def vqe_cost_2d(params_2d):
    # Use first two parameters, fix others
    full_params = np.random.uniform(0, 2*np.pi, 8)
    full_params[0] = params_2d[0]
    full_params[1] = params_2d[1]
    return vqe.cost_function(full_params)

plot_optimization_landscape_3d(
    vqe_cost_2d,
    [(0, 2*np.pi), (0, 2*np.pi)],
    'VQE Cost Landscape'
)

# Create simplified 2D cost landscape for QAOA
print("\nGenerating QAOA cost landscape...")
def qaoa_cost_2d(params_2d):
    return qaoa.cost_function(params_2d)

plot_optimization_landscape_3d(
    qaoa_cost_2d,
    [(0, np.pi), (0, np.pi)],
    'QAOA Cost Landscape'
)

# Summary statistics
print("\n" + "="*70)
print("SUMMARY STATISTICS")
print("="*70)

print("\nVQE Results:")
print(f"  Final Energy: {result_vqe.fun:.6f}")
print(f"  Convergence: {len(vqe.optimization_history)} iterations")
print(f"  Computation Time: {vqe_time:.2f}s")

print("\nQAOA Results:")
print(f"  Max Cut Value: {-result_qaoa.fun:.6f}")
print(f"  Convergence: {len(qaoa.optimization_history)} iterations")
print(f"  Computation Time: {qaoa_time:.2f}s")

print("\n" + "="*70)
print("EXECUTION COMPLETED SUCCESSFULLY")
print("="*70)