from qiskit_nature.units import DistanceUnit
from qiskit_nature.second_q.drivers import PySCFDriver

driver = PySCFDriver(
    atom="H 0 0 0; H 0 0 0.735",
    basis="sto3g",
    charge=0,
    spin=0,
    unit=DistanceUnit.ANGSTROM,
)

es_problem = driver.run()

from qiskit_nature.second_q.mappers import JordanWignerMapper

mapper = JordanWignerMapper()

from qiskit.algorithms.eigensolvers import NumPyEigensolver

numpy_solver = NumPyEigensolver(filter_criterion=es_problem.get_default_filter_criterion())

from qiskit.algorithms.minimum_eigensolvers import VQE
from qiskit.algorithms.optimizers import SLSQP
from qiskit.primitives import Estimator
from qiskit_nature.second_q.algorithms import GroundStateEigensolver, QEOM
from qiskit_nature.second_q.circuit.library import HartreeFock, UCCSD

ansatz = UCCSD(
    es_problem.num_spatial_orbitals,
    es_problem.num_particles,
    mapper,
    initial_state=HartreeFock(
        es_problem.num_spatial_orbitals,
        es_problem.num_particles,
        mapper,
    ),
)

estimator = Estimator()
# This first part sets the ground state solver
# see more about this part in the ground state calculation tutorial
solver = VQE(estimator, ansatz, SLSQP())
solver.initial_point = [0.0] * ansatz.num_parameters
gse = GroundStateEigensolver(mapper, solver)

# The qEOM algorithm is simply instantiated with the chosen ground state solver and Estimator primitive
qeom_excited_states_solver = QEOM(gse, estimator, "sd")

from qiskit_nature.second_q.algorithms import ExcitedStatesEigensolver

numpy_excited_states_solver = ExcitedStatesEigensolver(mapper, numpy_solver)
numpy_results = numpy_excited_states_solver.solve(es_problem)

qeom_results = qeom_excited_states_solver.solve(es_problem)

print(numpy_results)
print("\n\n")
print(qeom_results)

import numpy as np


def filter_criterion(eigenstate, eigenvalue, aux_values):
    return np.isclose(aux_values["ParticleNumber"][0], 2.0)


new_numpy_solver = NumPyEigensolver(filter_criterion=filter_criterion)
new_numpy_excited_states_solver = ExcitedStatesEigensolver(mapper, new_numpy_solver)
new_numpy_results = new_numpy_excited_states_solver.solve(es_problem)

print(new_numpy_results)

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright