import qiskit_alt
qiskit_alt.project.ensure_init()

# The following line would also initialize qiskit_alt behind the scenes, but we have made it explicity above.
# import qiskit_alt.electronic_structure

from qiskit.providers import aer  # make sure this is available

from qiskit_nature.drivers import UnitsType, 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.algorithms import VQE
from qiskit.circuit.library import TwoLocal
from qiskit.utils import QuantumInstance
from qiskit import Aer
from qiskit.algorithms.optimizers import COBYLA

import qiskit_alt, qiskit_alt.electronic_structure  # already done above, but this does not hurt
from qiskit_alt.pauli_operators import PauliSum_to_SparsePauliOp
from qiskit.opflow import PauliSumOp

geometry = [["H", [0.0, 0.0, 0.0]], ["H", [0.0, 0.0, 0.735]]]
basis="sto3g"

molecule = Molecule(
    geometry=geometry, charge=0, multiplicity=1
)
driver = ElectronicStructureMoleculeDriver(
    molecule, basis=basis, driver_type=ElectronicStructureDriverType.PYSCF
)

es_problem = ElectronicStructureProblem(driver)
second_q_op= es_problem.second_q_ops()
qubit_converter = QubitConverter(mapper=JordanWignerMapper())
qubit_op = qubit_converter.convert(second_q_op[0])
# Print the underlying data stored backing qubit_op
qubit_op.primitive.simplify()

# Compute the Fermionic operator of the molecule 
fermi_op = qiskit_alt.electronic_structure.fermionic_hamiltonian(geometry, basis)

# Convert the Fermionic operator to a Pauli operator using the Jordan-Wigner transform
pauli_op = qiskit_alt.electronic_structure.jordan_wigner(fermi_op)

type(pauli_op)

pauli_op.simplify()

(pauli_op - qubit_op.primitive).simplify()

fermi_op[1]  # In qiskit_alt, the Fermi operator carries the identity term

nuclear_energy_term = fermi_op[1].coeff
nuclear_energy_term

# Define the ansatz
ansatz = TwoLocal(
    rotation_blocks=["h", "rx"],
    entanglement_blocks="cz",
    entanglement="full",
    reps=2,
    parameter_prefix="y",
)

# Initialize the COBYLA optimizer
optimizer = COBYLA(maxiter=500)

# Select the backend for the quantum_instance
backend = QuantumInstance(Aer.get_backend('statevector_simulator'), 
                                        seed_simulator=100, 
                                        seed_transpiler=100)
# Run VQE algorithm
vqe = VQE(ansatz = ansatz, 
          optimizer = optimizer,
          quantum_instance = backend)

# Compute the ground-state energy of the molecule 
nature_result = vqe.compute_minimum_eigenvalue(operator=qubit_op)
nature_energy = nature_result.eigenvalue.real + nuclear_energy_term # Include constant term that is normally thrown away
print("The ground-state energy of the Hydrogen molecule is {} Hatree".format(round(nature_energy,3)))

# Convert the Pauli operator into a sum of Pauli operators, the required input format
pauli_sum_op = PauliSumOp(pauli_op)

# Compute the ground-state energy of the molecule
alt_result = vqe.compute_minimum_eigenvalue(operator=pauli_sum_op)
print("The ground-state energy of the Hydrogen molecule is {} Hatree".format(round(alt_result.eigenvalue.real,3)))

