from qiskit.quantum_info import SparsePauliOp

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

from qiskit.circuit.library import TwoLocal
from qiskit.algorithms.optimizers import SLSQP

ansatz = TwoLocal(3, rotation_blocks=["ry", "rz"], entanglement_blocks="cz", reps=1)

optimizer = SLSQP()
ansatz.decompose().draw('mpl')

from qiskit.primitives import Sampler, Estimator
from qiskit.algorithms.state_fidelities import ComputeUncompute

estimator = Estimator()
sampler = Sampler()
fidelity = ComputeUncompute(sampler)

k = 3
betas = [33, 33, 33]

counts = []
values = []
steps = []

def callback(eval_count, params, value, meta, step):
    counts.append(eval_count)
    values.append(value)
    steps.append(step)

from qiskit.algorithms.eigensolvers import VQD

vqd = VQD(estimator, fidelity, ansatz, optimizer, k=k, betas=betas, callback=callback)
result = vqd.compute_eigenvalues(operator = H2_op)
vqd_values = result.optimal_values

print(vqd_values)

import numpy as np
import pylab

pylab.rcParams["figure.figsize"] = (12, 8)

steps = np.asarray(steps)
counts = np.asarray(counts)
values = np.asarray(values)

for i in range(1,4):
    _counts = counts[np.where(steps == i)]
    _values = values[np.where(steps == i)]
    pylab.plot(_counts, _values, label=f"State {i-1}")
    
pylab.xlabel("Eval count")
pylab.ylabel("Energy")
pylab.title("Energy convergence for each computed state")
pylab.legend(loc="upper right");

from qiskit.algorithms.eigensolvers import NumPyEigensolver
from qiskit.opflow import PauliSumOp

exact_solver = NumPyEigensolver(k=3)
exact_result = exact_solver.compute_eigenvalues(PauliSumOp(H2_op))
ref_values = exact_result.eigenvalues

print(f"Reference values: {ref_values}")
print(f"VQD values: {vqd_values}")

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright