import numpy as np
import pylab
from qiskit_chemistry import QiskitChemistry

# Input dictionary to configure Qiskit Chemistry for the chemistry problem.
qiskit_chemistry_dict = {
    'driver': {'name': 'PSI4'},
    'PSI4': '',
    'algorithm': {'name': 'ExactEigensolver'},
}
# PSI4 config here is a multi-line string that we update using format()
# To do so all other curly brackets from PSI4 config must be doubled
psi4_cfg = """
molecule h2 {{
   0 1
   H 0.0 0.0 -{0}
   H 0.0 0.0 {0}
}}

set {{
  basis {1}
  scf_type pk
}}
"""
basis_sets = ['sto-3g', '3-21g', '6-31g']
start = 0.5  # Start distance
by    = 0.5  # How much to increase distance by
steps = 20   # Number of steps to increase by
energies  = np.empty([len(basis_sets), steps+1])
distances = 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
    for j in range(len(basis_sets)):
        qiskit_chemistry_dict['PSI4'] = psi4_cfg.format(d/2, basis_sets[j]) 
        solver = QiskitChemistry()
        result = solver.run(qiskit_chemistry_dict)
        energies[j][i] = result['energy']
    distances[i] = d
print(' --- complete')

print('Distances: ', distances)
print('Energies:', energies)

for j in range(len(basis_sets)):
    pylab.plot(distances, energies[j], label=basis_sets[j])
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('H2 Ground State Energy in different basis sets')
pylab.legend(loc='upper right')

