from qiskit.quantum_info.operators import SparsePauliOp

H2_op = SparsePauliOp(
    ["II", "IZ", "ZI", "ZZ", "XX"],
    coeffs=[
        -1.052373245772859,
        0.39793742484318045,
        -0.39793742484318045,
        -0.01128010425623538,
        0.18093119978423156,
    ],
)

aux_op1 = SparsePauliOp(["ZZ"], coeffs=[2.0])
aux_op2 = SparsePauliOp(["XX"], coeffs=[-2.0])


from IPython.display import clear_output
import matplotlib.pyplot as plt

counts = []
values = []

def plot_energy(eval_count, parameters, mean, metadata):
    counts.append(eval_count)
    values.append(mean)
    clear_output(wait=True)
    plt.plot(counts, values)
    plt.xlabel("Evaluation count")
    plt.ylabel("Energy")
    plt.title("Energy convergence")
    plt.show()

from qiskit.circuit.library import TwoLocal
from qiskit.primitives import Estimator
from qiskit.algorithms.minimum_eigensolvers import VQE, VQEResult
from qiskit.algorithms.optimizers import SLSQP

# Reset plotting arrays.
counts.clear()
values.clear()

# Define the central primitive.
estimator = Estimator()

# Define our ansatz 
ansatz = TwoLocal(rotation_blocks="ry", entanglement_blocks="cz")
slsqp = SLSQP()

# Instantiate VQE. Note that all arguments after the optimizer must be specified by keyword.
vqe = VQE(estimator, ansatz, slsqp, callback=plot_energy)

# Compute the minimum eigenvalue. Not the primary operator is the H2_op. The auxillary operators are just variants for illustration.
# Note that None values are no longer supported. Zero values will be interpreted as a zero operator.
result = vqe.compute_minimum_eigenvalue(H2_op, [aux_op1, aux_op2, 0])
print(result)

# Reset plotting arrays.
counts.clear()
values.clear()

result = vqe.compute_minimum_eigenvalue(H2_op, {"aux_op1": aux_op1, "aux_op2": aux_op2})
print(result)

from IPython.display import clear_output
import matplotlib.pyplot as plt
import numpy as np

counts = []
values = []
std_dev = []

def plot_energy_error(eval_count, parameters, mean, metadata):
    counts.append(eval_count)
    values.append(mean)
    variance = metadata.get("variance", 0.0)
    shots = metadata.get("shots", 0.0)
    std_dev.append(np.sqrt(variance / shots) if shots > 0 else 0.0)
    clear_output(wait=True)
    plt.errorbar(counts, values, std_dev)
    plt.xlabel("Evaluation count")
    plt.ylabel("Energy")
    plt.title("Energy convergence")
    plt.show()

counts.clear()
values.clear()

from qiskit.algorithms.optimizers import COBYLA

cobyla = COBYLA(maxiter=1000)

estimator = Estimator(options={"shots": 1024})

vqe = VQE(estimator, ansatz, cobyla, callback=plot_energy_error)
result = vqe.compute_minimum_eigenvalue(H2_op, [aux_op1, aux_op2])
print(result)

counts.clear()
values.clear()

from qiskit.algorithms.gradients import FiniteDiffEstimatorGradient

estimator = Estimator()

gradient = FiniteDiffEstimatorGradient(estimator, epsilon=0.001)

vqe = VQE(estimator, ansatz, slsqp, gradient=gradient, callback=plot_energy)
result = vqe.compute_minimum_eigenvalue(H2_op, [H2_op ** 2, H2_op / 2])
print(result)

counts.clear()
values.clear()

from qiskit.algorithms.gradients import ParamShiftEstimatorGradient
from qiskit.algorithms.optimizers import GradientDescent

estimator = Estimator()

gradient = FiniteDiffEstimatorGradient(estimator, epsilon=0.001)

vqe = VQE(estimator,
          ansatz,
          GradientDescent(maxiter=200, learning_rate=0.1),
          gradient=ParamShiftEstimatorGradient(estimator),
)

result = vqe.compute_minimum_eigenvalue(H2_op, [H2_op ** 2, H2_op / 2])
print(result)

from qiskit.circuit.library import TwoLocal
from qiskit.primitives import Sampler
from qiskit.algorithms.minimum_eigensolvers import SamplingVQE
from qiskit.algorithms.optimizers import SLSQP

# Reset plotting arrays.
counts.clear()
values.clear()

# Define the central primitive.
sampler = Sampler()

# Instantiate SamplingVQE. Note that all arguments after the optimizer must be specified by keyword.
sampling_vqe = SamplingVQE(sampler, ansatz, slsqp, callback=plot_energy)

result = sampling_vqe.compute_minimum_eigenvalue(H2_op)

# Reset plotting arrays.
counts.clear()
values.clear()

diag_op = SparsePauliOp(["ZZ", "IZ", "II"], coeffs=[1, -0.5, 0.12])

sampler = Sampler()
from qiskit.circuit.library import RealAmplitudes

sampling_vqe = SamplingVQE(
            sampler,
            RealAmplitudes(2, reps=1),
            SLSQP(),
            aggregation=0.1,
            callback=plot_energy,
        )

sampling_vqe.compute_minimum_eigenvalue(operator=diag_op)

