import numpy as np
from qiskit.algorithms import VQE
from qiskit.algorithms.optimizers import GradientDescent
from qiskit.circuit.library import EfficientSU2
from qiskit.opflow import Z, I, Gradient, AerPauliExpectation
from qiskit.providers.aer import AerSimulator

hamiltonian = (Z ^ Z ^ I) + (Z ^ I ^ I)
ansatz = EfficientSU2(3, reps=1, entanglement="linear")

ansatz.decompose().draw("mpl", style="iqx")

gradient = Gradient(method="param_shift")
backend = AerSimulator()

energies = []
def store_energy(*args):
    energies.append(args[2])

vqe = VQE(ansatz, 
          optimizer=GradientDescent(maxiter=300, learning_rate=0.01, callback=store_energy),
          gradient=gradient,
          expectation=AerPauliExpectation(),
          initial_point=np.arange(ansatz.num_parameters),
          quantum_instance=backend)
result = vqe.compute_minimum_eigenvalue(hamiltonian)

import matplotlib.pyplot as plt 

plt.figure(figsize=(10, 6))
plt.plot(energies, color="royalblue", label="VQE")
plt.axhline(y=-2, ls="--", color="crimson", label="target")
plt.xlabel("iteration")
plt.ylabel("energy");

from time import time
from qiskit.opflow import CircuitSampler, StateFn

num_qubits = 5
gradient = Gradient(method="param_shift")
sampler = CircuitSampler(backend)
hamiltonian = Z ^ 5
timings = []

for reps in range(10):
    circuit = EfficientSU2(num_qubits, reps=reps, entanglement="pairwise")
    value_dict = dict(zip(circuit.parameters, np.arange(circuit.num_parameters)))
    expectation = StateFn(hamiltonian, is_measurement=True) @ StateFn(circuit)
    start = time()
    _ = sampler.convert(gradient.convert(expectation), params=value_dict).eval()
    timings.append(time() - start)

plt.figure(figsize=(10, 6))
plt.plot(10 * (1 + np.arange(10)), timings)
plt.xlabel("number of parameters")
plt.ylabel("time $s$");

from surfer.gradient import ForwardGradient  # classical O(d^2) method
from surfer.gradient import ReverseGradient  # classical O(d) method!

reverse = ReverseGradient()

hamiltonian = (Z ^ Z ^ I) + (Z ^ I ^ I)
circuit = EfficientSU2(3, reps=1, entanglement="linear")
values = np.random.random(ansatz.num_parameters)

grad = reverse.compute(hamiltonian, circuit, values)
print(grad)

from surfer.gradient import ForwardGradient

reverse_timings = []
forward_timings = []
forward = ForwardGradient()

hamiltonian = Z ^ 5

for reps in range(10):
    circuit = EfficientSU2(num_qubits, reps=reps, entanglement="pairwise")
    values = np.arange(circuit.num_parameters)
    start = time()
    _ = reverse.compute(hamiltonian, circuit, values)
    reverse_timings.append(time() - start)
    
    circuit = circuit.decompose()
    start = time()
    _ = forward.compute(hamiltonian, circuit, values)
    forward_timings.append(time() - start)

plt.figure(figsize=(10, 6))
plt.plot(10 * (1 + np.arange(10)), timings, label="Qiskit")
plt.plot(10 * (1 + np.arange(10)), reverse_timings, label="Reverse gradients")
plt.plot(10 * (1 + np.arange(10)), forward_timings, label="Forward gradients")
plt.xlabel("number of parameters")
plt.ylabel("time $s$")
plt.legend();

from surfer.qfi import ReverseQFI

reverse_qfi = ReverseQFI()

circuit = EfficientSU2(3, reps=1, entanglement="linear")
values = np.random.random(ansatz.num_parameters)

qfi = reverse_qfi.compute(circuit, values)

from qiskit.opflow import QFI

num_qubits = 5

qfi = QFI()
qfi_timings = []
sampler = CircuitSampler(backend)

for reps in [0, 1, 2, 3, 4, 8]:
    circuit = EfficientSU2(num_qubits, reps=reps, entanglement="pairwise")
    value_dict = dict(zip(circuit.parameters, np.arange(circuit.num_parameters)))

    start = time()
    _ = sampler.convert(qfi.convert(StateFn(circuit)), params=value_dict).eval()
    qfi_timings.append(time() - start)

num_qubits = 5
reverse_qfi_timings = []

for reps in 2 ** np.arange(6):
    circuit = EfficientSU2(num_qubits, reps=reps, entanglement="pairwise")
    values = np.arange(circuit.num_parameters)

    start = time()
    _ = reverse_qfi.compute(circuit, values)
    reverse_qfi_timings.append(time() - start)

plt.figure(figsize=(10, 6))
plt.plot(10 * (1 + np.array([0, 1, 2, 3, 4, 8])), qfi_timings, label="Qiskit")
plt.plot(10 * (1 + 2 ** np.arange(6)), reverse_qfi_timings, label="Classically efficient QFI")
plt.xlabel("number of parameters")
plt.ylabel("time $s$")
plt.legend();

