from qiskit import Aer
from qiskit.utils import QuantumInstance

from qiskit_nature.algorithms import VQEUCCFactory
from qiskit_nature.algorithms.initial_points.mp2_initial_point import MP2InitialPoint
from qiskit_nature.drivers import Molecule
from qiskit_nature.drivers.second_quantization import (
    ElectronicStructureDriverType,
    ElectronicStructureMoleculeDriver,
)
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_nature.algorithms import GroundStateEigensolver

from qiskit_nature.transformers.second_quantization.electronic import FreezeCoreTransformer

molecule = Molecule(
    geometry=[["H", [0.0, 0.0, 0.0]], ["H", [0.0, 0.0, 0.735]]], charge=0, multiplicity=1
)

driver = ElectronicStructureMoleculeDriver(
    molecule, basis="sto3g", driver_type=ElectronicStructureDriverType.PYSCF
)

freeze_core = FreezeCoreTransformer()

problem = ElectronicStructureProblem(driver, transformers=[freeze_core])

qubit_converter = QubitConverter(JordanWignerMapper())

initial_point = MP2InitialPoint()

quantum_instance = QuantumInstance(backend=Aer.get_backend("aer_simulator_statevector"))
vqe_ucc_factory = VQEUCCFactory(quantum_instance, initial_point=initial_point)

vqe_solver = GroundStateEigensolver(qubit_converter, vqe_ucc_factory)
res = vqe_solver.solve(problem)

print(res)

from qiskit.algorithms.optimizers import L_BFGS_B
from qiskit.algorithms import VQE

from qiskit_nature.circuit.library import UCCSD
from qiskit_nature.settings import settings
from qiskit_nature.drivers import Molecule


settings.dict_aux_operators = True

molecule = Molecule(
    geometry=[["H", [0.0, 0.0, 0.0]], ["H", [0.0, 0.0, 0.735]]],
    charge=0,
    multiplicity=1,
)

driver = ElectronicStructureMoleculeDriver(
    molecule, basis="sto3g", driver_type=ElectronicStructureDriverType.PYSCF
)

problem = ElectronicStructureProblem(driver)

# generate the second-quantized operators
second_q_ops = problem.second_q_ops()
main_op = second_q_ops["ElectronicEnergy"]

driver_result = problem.grouped_property_transformed

particle_number = driver_result.get_property("ParticleNumber")
electronic_energy = driver_result.get_property("ElectronicEnergy")

num_particles = (particle_number.num_alpha, particle_number.num_beta)
num_spin_orbitals = particle_number.num_spin_orbitals

# setup the classical optimizer for VQE
optimizer = L_BFGS_B()

# setup the qubit converter
converter = QubitConverter(JordanWignerMapper())

# map to qubit operators
qubit_op = converter.convert(main_op, num_particles=num_particles)

# setup the ansatz for VQE
ansatz = UCCSD(
    qubit_converter=converter,
    num_particles=num_particles,
    num_spin_orbitals=num_spin_orbitals,
)

mp2_initial_point = MP2InitialPoint()
mp2_initial_point.grouped_property = driver_result
mp2_initial_point.ansatz = ansatz
initial_point = mp2_initial_point.to_numpy_array()

# set the backend for the quantum computation
backend = Aer.get_backend("aer_simulator_statevector")

# setup and run VQE
algorithm = VQE(
    ansatz, optimizer=optimizer, quantum_instance=backend, initial_point=initial_point
)

result = algorithm.compute_minimum_eigenvalue(qubit_op)
print(result.eigenvalue.real)

electronic_structure_result = problem.interpret(result)
print(electronic_structure_result)
