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', 'transformation': 'full', 
                 'qubit_mapping': 'parity', 'two_qubit_reduction': True},
    'algorithm': {'name': 'VQE'},
    'optimizer': {'name': 'SPSA', 'max_trials': 350},
    'variational_form': {'name': 'RYRZ', 'depth': 3, 'entanglement': 'full'}
}
molecule = 'H .0 .0 -{0}; H .0 .0 {0}'
algorithms = ['VQE', 'ExactEigensolver']
backends = [{'name': 'qasm_simulator', 'shots': 1024},
             None
           ]

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(algorithms), steps+1])
hf_energies = np.empty(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) 
    for j in range(len(algorithms)):
        qiskit_chemistry_dict['algorithm']['name'] = algorithms[j]
        if backends[j] is not None:
            qiskit_chemistry_dict['backend'] = backends[j]
        else:
            qiskit_chemistry_dict.pop('backend')
        solver = QiskitChemistry()
        result = solver.run(qiskit_chemistry_dict)
        energies[j][i] = result['energy']
        hf_energies[i] = result['hf_energy']
    distances[i] = d
print(' --- complete')

print('Distances: ', distances)
print('Energies:', energies)
print('Hartree-Fock energies:', hf_energies)

pylab.plot(distances, hf_energies, label='Hartree-Fock')
for j in range(len(algorithms)):
    pylab.plot(distances, energies[j], label=algorithms[j])
pylab.xlabel('Interatomic distance (Angstrom)')
pylab.ylabel('Energy (Hartree)')
pylab.title('H2 Ground State Energy')
pylab.legend(loc='upper right')

pylab.plot(distances, np.subtract(hf_energies, energies[1]), label='Hartree-Fock')
pylab.plot(distances, np.subtract(energies[0], energies[1]), label=algorithms[0])
pylab.xlabel('Interatomic distance (Angstrom)')
pylab.ylabel('Energy (Hartree)')
pylab.title('Energy difference from ExactEigensolver')
pylab.legend(loc='upper left')

