from collections import OrderedDict
from qiskit import LegacySimulators
from qiskit.transpiler import PassManager
from qiskit_aqua import AquaError
from qiskit_aqua import QuantumInstance
from qiskit_aqua.algorithms import ExactEigensolver
from qiskit_aqua.algorithms import QPE
from qiskit_aqua.components.iqfts import Standard
from qiskit_chemistry import FermionicOperator
from qiskit_chemistry import QiskitChemistry
from qiskit_chemistry.drivers import ConfigurationManager
from qiskit_chemistry.aqua_extensions.components.initial_states import HartreeFock
import time

distance = 0.735
cfg_mgr = ConfigurationManager()
pyscf_cfg = OrderedDict([
    ('atom', 'H .0 .0 .0; H .0 .0 {}'.format(distance)),
    ('unit', 'Angstrom'),
    ('charge', 0),
    ('spin', 0),
    ('basis', 'sto3g')
])
section = {}
section['properties'] = pyscf_cfg
try:
    driver = cfg_mgr.get_driver_instance('PYSCF')
except ModuleNotFoundError:
    raise AquaError('PYSCF driver does not appear to be installed')

molecule = driver.run(section)
qubit_mapping = 'parity'
fer_op = FermionicOperator(h1=molecule.one_body_integrals, h2=molecule.two_body_integrals)
qubit_op = fer_op.mapping(map_type=qubit_mapping,threshold=1e-10).two_qubit_reduced_operator(2)

exact_eigensolver = ExactEigensolver(qubit_op, k=1)
result_ee = exact_eigensolver.run()
reference_energy = result_ee['energy']
print('The exact ground state energy is: {}'.format(result_ee['energy']))

num_particles = molecule.num_alpha + molecule.num_beta
two_qubit_reduction = True
num_orbitals = qubit_op.num_qubits + (2 if two_qubit_reduction else 0)

num_time_slices = 50
n_ancillae = 9

state_in = HartreeFock(qubit_op.num_qubits, num_orbitals,
                       num_particles, qubit_mapping, two_qubit_reduction)
iqft = Standard(n_ancillae)

qpe = QPE(qubit_op, state_in, iqft, num_time_slices, n_ancillae,
          paulis_grouping='random', expansion_mode='suzuki',
          expansion_order=2, shallow_circuit_concat=True)
backend = LegacySimulators.get_backend('qasm_simulator')
quantum_instance = QuantumInstance(backend, shots=100, pass_manager=PassManager())
result_qpe = qpe.run(quantum_instance)
print('The ground state energy as computed by QPE is: {}'.format(result_qpe['energy']))

molecule = 'H .0 .0 0; H .0 .0 {}'.format(distance)

# Input dictionary to configure Qiskit Chemistry for the chemistry problem.
qiskit_chemistry_qpe_dict = {
    'driver': {'name': 'PYSCF'},
    'PYSCF': {
        'atom': molecule, 
        'basis': 'sto3g'
    },
    'operator': {'name': 'hamiltonian', 'transformation': 'full', 'qubit_mapping': 'parity'},
    'algorithm': {
        'name': 'QPE',
        'num_ancillae': 9,
        'num_time_slices': 50,
        'expansion_mode': 'suzuki',
        'expansion_order': 2,
    },
    'initial_state': {'name': 'HartreeFock'},
    'backend': {'shots': 100}
}

qiskit_chemistry_ees_dict = {
    'driver': {'name': 'PYSCF'},
    'PYSCF': {'atom': molecule, 'basis': 'sto3g'},
    'operator': {'name': 'hamiltonian', 'transformation': 'full', 'qubit_mapping': 'parity'},
    'algorithm': {
        'name': 'ExactEigensolver',
    }
}

result_qpe = QiskitChemistry().run(qiskit_chemistry_qpe_dict, backend=backend)
result_ees = QiskitChemistry().run(qiskit_chemistry_ees_dict)

print('The groundtruth total ground state energy is           {}.'.format(
    result_ees['energy'] - result_ees['nuclear_repulsion_energy']
))
print('The total ground state energy as computed by QPE is    {}.'.format(
    result_qpe['energy'] - result_qpe['nuclear_repulsion_energy']
))
print('In comparison, the Hartree-Fock ground state energy is {}.'.format(
    result_ees['hf_energy'] - result_ees['nuclear_repulsion_energy']
))