#!/usr/bin/env python3

# Simple example where we train a parameterized quantum circuit
# to minimize the expectation value of a Pauli-Z operator.
import pennylane as qml
from pennylane import numpy as np

# Define a quantum device with one qubit
dev = qml.device("default.qubit", wires=1)

# Define a variational quantum circuit
@qml.qnode(dev)
def circuit(params):
    qml.RX(params[0], wires=0)  # Apply a rotation around the X-axis
    qml.RY(params[1], wires=0)  # Apply a rotation around the Y-axis
    return qml.expval(qml.PauliZ(0))

# Define the cost function to minimize
def cost(params):
    return circuit(params)

# Initialize parameters
init_params = np.array([0.01, 0.01], requires_grad=True)

# Set up the optimizer
opt = qml.GradientDescentOptimizer(stepsize=0.1)

# Number of optimization steps
steps = 100

# Optimize the circuit parameters
params = init_params
for i in range(steps):
    params = opt.step(cost, params)
    if (i + 1) % 10 == 0:
        print(f"Step {i + 1}: Cost = {cost(params):.6f}, Params = {params}")

print(f"Optimized Parameters: {params}")
print(f"Final Cost: {cost(params)}")