from qiskit_nature.drivers import PySCFDriver
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import JordanWignerMapper
from qiskit.algorithms import VQE, NumPyMinimumEigensolver
from qiskit.algorithms.optimizers import SPSA
from qiskit.circuit.library import TwoLocal
from qiskit import Aer
from qiskit.utils import QuantumInstance

def run_vqe_h2():
    """Runs the VQE algorithm to find the ground state energy of the H2 molecule."""
    # 1. Define the molecule (H2 at its equilibrium bond length)
    driver = PySCFDriver(atom='H .0 .0 .0; H .0 .0 0.735')

    # 2. Define the electronic structure problem
    problem = ElectronicStructureProblem(driver)

    # 3. Define the qubit converter
    qubit_converter = QubitConverter(mapper=JordanWignerMapper())

    # 4. Set up the VQE algorithm
    # Choose an optimizer
    optimizer = SPSA(maxiter=100)

    # Choose an ansatz (trial wavefunction)
    ansatz = TwoLocal(rotation_blocks='ry', entanglement_blocks='cz')

    # Choose a quantum backend (simulator)
    backend = Aer.get_backend('statevector_simulator')
    quantum_instance = QuantumInstance(backend)

    # Create the VQE instance
    vqe = VQE(ansatz=ansatz, optimizer=optimizer, quantum_instance=quantum_instance)

    # 5. Run the VQE algorithm
    # First, get the qubit operator from the problem
    qubit_op = qubit_converter.convert(problem.second_q_ops()[0])[0]
    result = vqe.compute_minimum_eigenvalue(operator=qubit_op)

    print(f"VQE Result: {result.eigenvalue.real}")

    # For comparison, we can also find the exact ground state energy using a classical solver
    np_solver = NumPyMinimumEigensolver()
    exact_result = np_solver.compute_minimum_eigenvalue(operator=qubit_op)
    print(f"Exact Result: {exact_result.eigenvalue.real}")

if __name__ == "__main__":
    run_vqe_h2()