from qiskit_nature.drivers import Molecule
from qiskit_nature.drivers.second_quantization import ElectronicStructureDriverType, ElectronicStructureMoleculeDriver

dist = 0.72
molecule = Molecule(geometry=[['H', [0., 0., 0.]], ['H', [0., 0., dist]]])
driver = ElectronicStructureMoleculeDriver(molecule, basis='sto3g', \
                                           driver_type=ElectronicStructureDriverType.PYSCF)

print(molecule)

from qiskit_nature.problems.second_quantization import ElectronicStructureProblem

es_problem = ElectronicStructureProblem(driver)
second_q_ops = es_problem.second_q_ops()
print(second_q_ops[0])

from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import ParityMapper

qubit_converter = QubitConverter(ParityMapper(), two_qubit_reduction=True)
hamiltonian = qubit_converter.convert(second_q_ops[0], num_particles=es_problem.num_particles)
print(hamiltonian)

from qiskit.algorithms import NumPyMinimumEigensolver
from qiskit.algorithms import MinimumEigensolverResult

# Generate the exact solution
sol = MinimumEigensolverResult()
sol.eigenvalue = NumPyMinimumEigensolver().compute_minimum_eigenvalue(hamiltonian).eigenvalue
exact_solution = es_problem.interpret(sol).total_energies[0]

print('exact_solution:\n{}'.format(exact_solution))

from qiskit.opflow import TaperedPauliSumOp
import numpy as np

ops = []
h = hamiltonian.primitive
h.coeffs = np.real(hamiltonian.coeffs)
ops.append(TaperedPauliSumOp(h, hamiltonian.z2_symmetries))

from qiskit_ibm_runtime import QiskitRuntimeService, Estimator
service = QiskitRuntimeService()
options = {"backend" : "ibmq_qasm_simulator"}

from qiskit.algorithms.optimizers import NELDER_MEAD
from qiskit.circuit.library import RealAmplitudes
circuit = RealAmplitudes(num_qubits=hamiltonian.num_qubits, reps=2)

circuit.decompose().draw()

convergence = []

with Estimator(
    circuits=circuit,
    observables=ops,
    options=options,
    service=service
) as estimator:
    
    def evaluate_expectation(x):
        x = list(x)
        results = estimator(circuits=[0], observables=[0], parameter_values=[x]).values[0]
        return np.real(results)

    def callback(fx):
        expval = evaluate_expectation(fx)
        print('expval: {}'.format(expval))
        convergence.append(expval)

    np.random.seed(1)
    initial_point = np.random.uniform(-np.pi, np.pi, len(circuit.parameters))
    optimizer = NELDER_MEAD(40, callback=callback)
    result = optimizer.minimize(fun=evaluate_expectation, x0=initial_point)


total_energy = []
for i in range(len(convergence)):
    sol = MinimumEigensolverResult()
    sol.eigenvalue = convergence[i]
    sol = es_problem.interpret(sol).total_energies[0]
    total_energy.append(sol)

print(total_energy)

import matplotlib.pyplot as plt
plt.rcParams["font.size"] = 14

plt.figure(figsize=(12, 6))
plt.plot(total_energy, label="Estimator VQE")
plt.axhline(y=exact_solution, color="tab:red", ls="--", label="Target")

plt.legend(loc="best")
plt.xlabel("Iteration")
plt.ylabel("Energy [H]")
plt.title("VQE energy")
plt.show()

import qiskit_ibm_runtime
qiskit_ibm_runtime.version.get_version_info()

from qiskit.tools.jupyter import *

%qiskit_version_table
%qiskit_copyright