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': 'PYSCF'},
    'PYSCF': {'atom': '', 'basis': 'sto3g'},
    'operator': {'name': 'hamiltonian', 'qubit_mapping': 'parity', 'two_qubit_reduction': True},
    'algorithm': {'name': 'ExactEigensolver', 'k': 4},
}
molecule = 'H .0 .0 -{0}; H .0 .0 {0}'

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([4, 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
    qiskit_chemistry_dict['PYSCF']['atom'] = molecule.format(d/2) 
    solver = QiskitChemistry()
    result = solver.run(qiskit_chemistry_dict)
    energies[:, i] = result['energies']
    distances[i] = d
print(' --- complete')

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

pylab.rcParams['figure.figsize'] = (12, 8)
for j in range(energies.shape[0]):
    label = 'Ground state' if j ==0 else 'Excited state {}'.format(j)
    pylab.plot(distances, energies[j], label=label)
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('H2 Ground and Excited States')
pylab.legend(loc='upper right')
pylab.show()

pylab.rcParams['figure.figsize'] = (6, 4)
prop_cycle = pylab.rcParams['axes.prop_cycle']
colors = prop_cycle.by_key()['color']
for j in range(energies.shape[0]):
    label = 'Ground state' if j ==0 else 'Excited state {}'.format(j)
    pylab.plot(distances, energies[j], color=colors[j], label=label)
    pylab.xlabel('Interatomic distance')
    pylab.ylabel('Energy')
    pylab.title('H2 {}'.format(label))
    pylab.legend(loc='upper right')
    pylab.show()