from qiskit.primitives import Estimator
from qiskit.algorithms.gradients import ParamShiftEstimatorGradient

estimator = Estimator()
ps = ParamShiftEstimatorGradient(estimator)

from qiskit.circuit.library import EfficientSU2
from qiskit.quantum_info import SparsePauliOp

def ising_hamiltonian(num_qubits):
    return SparsePauliOp.from_sparse_list(
        [("ZZ", [i, i + 1], 1) for i in range(num_qubits - 1)] + [("X", [i], -1) for i in range(num_qubits)],
        num_qubits=num_qubits
    )


num_qubits = 3
hamiltonian = ising_hamiltonian(num_qubits)
print(hamiltonian)

circuit = EfficientSU2(num_qubits, reps=1).decompose()
print(circuit.draw())

import numpy as np

values = np.random.random(circuit.num_parameters)

gradient = ps.run([circuit], [hamiltonian], [values]).result()
print(gradient.gradients[0])

from qiskit.algorithms.gradients import ReverseEstimatorGradient

reverse = ReverseEstimatorGradient()  # classical calculation, not primitive-based!

gradient = reverse.run([circuit], [hamiltonian], [values]).result()
print(gradient.gradients[0])

import time
import numpy as np
from qiskit.circuit.library import EfficientSU2

def time_gradients(num_layers, gradient, num_qubits=5, averages=5):
    circuit = EfficientSU2(num_qubits, reps=num_layers).decompose()
    hamiltonian = ising_hamiltonian(num_qubits)
    values = np.linspace(1, 2, num=circuit.num_parameters)
    
    timings = []
    for _ in range(averages):
        start = time.time()
        __ = gradient.run([circuit], [hamiltonian], [values]).result()
        timings.append(time.time() - start)
        
    return np.mean(timings), np.std(timings)

layers = [2, 4, 6, 8, 10]
num_qubits = 5

ps_times = np.array([time_gradients(num_layers, ps) for num_layers in layers])

rev_times = np.array([time_gradients(num_layers, reverse) for num_layers in layers])

import matplotlib.pyplot as plt

num_parameters = (np.array(layers) + 1) * 2 * num_qubits

plt.figure(figsize=(12, 6))
plt.plot(num_parameters, ps_times[:, 0], color="royalblue", label="circuit-based")
plt.fill_between(num_parameters, ps_times[:, 0] - ps_times[:, 1], ps_times[:, 0] + ps_times[:, 1], color="royalblue", alpha=0.2)
plt.plot(num_parameters, rev_times[:, 0], color="seagreen", label="reverse")
plt.fill_between(num_parameters, rev_times[:, 0] - rev_times[:, 1], rev_times[:, 0] + rev_times[:, 1], color="seagreen", alpha=0.2)
plt.title("Gradient benchmark")
plt.xlabel("number of parameters")
plt.ylabel("runtime in seconds");

from qiskit.algorithms.minimum_eigensolvers import VQE
from qiskit.algorithms.optimizers import GradientDescent
from qiskit.algorithms.time_evolvers.variational import RealMcLachlanPrinciple

def time_vqe(num_layers, gradient, num_qubits=4):
    circuit = EfficientSU2(num_qubits, reps=num_layers).decompose()
    hamiltonian = ising_hamiltonian(num_qubits)
    
    optimizer = GradientDescent(maxiter=100)
    initial_point = np.zeros(circuit.num_parameters)
    vqe = VQE(estimator, circuit, optimizer, gradient=gradient, initial_point=initial_point)
    
    start = time.time()
    _ = vqe.compute_minimum_eigenvalue(hamiltonian)
    took = time.time() - start
    print("done in", took)
    return took

num_qubits = 4
layers = [2, 4, 6, 8, 10, 12]
num_parameters = (np.array(layers) + 1) * 2 * num_qubits

vqe_ps_times = np.array([time_vqe(num_layers, ps) for num_layers in layers])
# vqe_ps_times = np.load("vqe_ps_times.npy")

vqe_rev_times = np.array([time_vqe(num_layers, reverse) for num_layers in layers])
# vqe_rev_times = np.load("vqe_rev_times.npy")

plt.figure(figsize=(12, 6))
plt.plot(num_parameters, vqe_ps_times / 60, color="royalblue", marker="o", label="circuit-based")
plt.plot(num_parameters, vqe_rev_times / 60, color="seagreen", marker="v", label="reverse")
plt.legend(loc="best")
plt.title("VQE with 100 steps")
plt.xlabel("number of parameters")
plt.ylabel("time in minutes");

from qiskit.algorithms.gradients import ReverseQGT

from qiskit.algorithms.time_evolvers import VarQRTE, TimeEvolutionProblem
from qiskit.algorithms.time_evolvers.variational import RealMcLachlanPrinciple

def time_varqte(num_layers, estimator, gradient, qgt, num_qubits=5, averages=1):
    circuit = EfficientSU2(num_qubits, reps=num_layers).decompose()
    hamiltonian = ising_hamiltonian(num_qubits)
    problem = TimeEvolutionProblem(hamiltonian, time=1)#, aux_operators=[hamiltonian])
    
    principle = RealMcLachlanPrinciple(qgt, gradient)
    initial_values = np.zeros(circuit.num_parameters)
    varqte = VarQRTE(circuit, initial_values, principle, estimator=estimator, num_timesteps=100)
    
    start = time.time()
    _ = varqte.evolve(problem)
    took = time.time() - start
    print("done in", took)
        
    return took

from qiskit.algorithms.gradients import LinCombEstimatorGradient, LinCombQGT

num_qubits = 4
layers = [2, 4, 6, 8]
num_parameters = (np.array(layers) + 1) * 2 * num_qubits

lcu_qgt = LinCombQGT(estimator)
lcu = LinCombEstimatorGradient(estimator)

# varqte_times = np.array([time_varqte(num_layers, estimator, lcu, lcu_qgt) for num_layers in layers])
varqte_times = np.load("varqte_times.npy")

rev_qgt = ReverseQGT()

# revqte_times = np.array([time_varqte(num_layers, estimator, reverse, rev_qgt) for num_layers in layers])
revqte_times = np.load("revqte_times.npy")

plt.figure(figsize=(12, 6))
plt.plot(num_parameters[:len(varqte_times)], varqte_times / 60, color="royalblue", marker="o", label="circuit-based")
plt.plot(num_parameters, revqte_times / 60, color="seagreen", marker="v", label="reverse")
plt.legend(loc="best")
plt.title("VarQTE with 100 timesteps")
plt.xlabel("number of parameters")
plt.ylabel("time in minutes");

