import pennylane as qml
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

# Enable 64-bit precision (crucial for accurately measuring vanishing gradients)
jax.config.update("jax_enable_x64", True)

# ==========================================
# 1 & 2. Circuit and Objective Definition (JAX interface)
# ==========================================

def create_objective(n_qubits, n_layers):
    dev = qml.device("default.qubit", wires=n_qubits)

    # Specify "jax" as the interface
    @qml.qnode(dev, interface="jax")
    def circuit(params):
        for l in range(n_layers):
            for i in range(n_qubits):
                qml.RX(params[l, i, 0], wires=i)
                qml.RZ(params[l, i, 1], wires=i)
            
            for i in range(n_qubits - 1):
                qml.CNOT(wires=[i, i + 1])
                
        return qml.expval(qml.PauliZ(0))

    return circuit

# ==========================================
# 3. JAX Barren Plateau Experiment Array
# ==========================================

qubit_counts = [8]
n_layers = 12  # Keeping original fixed depth logic intact
num_samples = 500  # Keeping original sample size

mean_hessians = []
var_hessians = []

print("Starting Fast JAX Barren Plateau Analysis (Original Logic)...")

# Initialize JAX random key
key = jax.random.PRNGKey(30)

for n_qubits in qubit_counts:
    print(f"\nAnalyzing {n_qubits} qubits...")
    
    circuit_fn = create_objective(n_qubits, n_layers)
    
    # JAX Magic: Compile the batched Hessian calculation
    batched_hessian_fn = jax.jit(jax.vmap(jax.hessian(circuit_fn)))
    
    # Generate all random parameters at once (Shape: 100, layers, qubits, 2)
    key, subkey = jax.random.split(key)
    params = jax.random.uniform(subkey, shape=(num_samples, n_layers, n_qubits, 2), minval=0, maxval=2*jnp.pi)
    
    # Compute all 100 Hessians instantly in one vectorized pass
    H_batched = batched_hessian_fn(params)
    
    # KEEPING ORIGINAL LOGIC: Flatten everything into one massive array
    all_hessian_elements = H_batched.flatten()
    
    mean_h = jnp.mean(all_hessian_elements)
    var_h = jnp.var(all_hessian_elements)
    
    mean_hessians.append(mean_h)
    var_hessians.append(var_h)
    
    print(f"Result for {n_qubits} qubits -> Mean: {mean_h:.8f}, Variance: {var_h:.8f}")

# ==========================================
# 4. Plotting the Results
# ==========================================

plt.figure(figsize=(8, 5))
plt.semilogy(qubit_counts, var_hessians, marker='o', linestyle='-', color='b')

plt.title("Hessian Variance vs. Number of Qubits (JAX, Fixed Layers)", fontsize=14)
plt.xlabel("Number of Qubits", fontsize=12)
plt.ylabel("Variance of Hessian Elements (Log Scale)", fontsize=12)
plt.xticks(qubit_counts)
plt.grid(True, which="both", ls="--", alpha=0.6)

plt.tight_layout()
plt.show()