import numpy as np
import pylab
from qiskit_chemistry import QiskitChemistry

# Input dictionary to configure Qiskit Chemistry for the chemistry problem.
# Note: In order to allow this to run reasonably quickly it takes advantage
#       of the ability to freeze core orbitals and remove unoccupied virtual
#       orbitals to reduce the size of the problem. The result without this
#       will be more accurate but it takes rather longer to run.
qiskit_chemistry_dict = {
    'driver': {'name': 'PYSCF'},
    'PYSCF': {'atom': '', 'basis': 'sto3g'},
    'algorithm': {'name': 'ExactEigensolver'},
    'operator': {'name': 'hamiltonian', 'freeze_core': True, 'orbital_reduction': [-3, -2]},
}
molecule = 'Li .0 .0 -{0}; H .0 .0 {0}'

start = 1.25 # Start distance
by    = 0.5  # How much to increase distance by
steps = 20   # Number of steps to increase by
energies  = np.empty(steps+1)
distances = np.empty(steps+1)
dipoles   = np.empty(steps+1)

print('Processing step __', end='')
for i in range(steps+1):
    print('\b\b{:2d}'.format(i), end='', flush=True)
    d = start + i*by/steps
    qiskit_chemistry_dict['PYSCF']['atom'] = molecule.format(d/2) 
    solver = QiskitChemistry()
    result = solver.run(qiskit_chemistry_dict)
    distances[i] = d
    energies[i] = result['energy']
    dipoles[i]  = result['total_dipole_moment']
print(' --- complete')

print('Distances: ', distances)
print('Energies:', energies)
print('Dipole moments:', dipoles)

pylab.plot(distances, energies)
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('LiH Ground State Energy')

pylab.plot(distances, dipoles)
pylab.xlabel('Interatomic distance')
pylab.ylabel('Moment a.u')
pylab.title('LiH Dipole Moment')

