import numpy as np
import pylab
from qiskit_chemistry import QiskitChemistry

# Input dictionary to configure Qiskit Chemistry for the chemistry problem.
qiskit_chemistry_dict = {
    'problem': {'random_seed': 50},
    'driver': {'name': 'PYSCF'},
    'PYSCF': {'atom': '', 'basis': 'sto3g'},
    'operator': {'name': 'hamiltonian', 'qubit_mapping': '', 'two_qubit_reduction': False},
    'algorithm': {'name': ''},
    'optimizer': {'name': 'L_BFGS_B', 'maxfun': 2500},
    'variational_form': {'name': 'RYRZ', 'depth': 5}
}
molecule = 'H .0 .0 -{0}; H .0 .0 {0}'

algorithms = ['VQE', 'ExactEigensolver']
mappings   = ['jordan_wigner', 'parity', 'bravyi_kitaev']
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(mappings), 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] 
        for k in range(len(mappings)):
            qiskit_chemistry_dict['operator']['qubit_mapping'] = mappings[k] 
            solver = QiskitChemistry()
            result = solver.run(qiskit_chemistry_dict)
            energies[k][j][i] = result['energy']
            hf_energies[i] = result['hf_energy']  # Independent of algorithm & mapping
    distances[i] = d
print(' --- complete')

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


pylab.rcParams['figure.figsize'] = (12, 8)
pylab.ylim(-1.14, -1.04)
pylab.plot(distances, hf_energies, label='Hartree-Fock')
for j in range(len(algorithms)):
    for k in range(len(mappings)):
        pylab.plot(distances, energies[k][j], label=algorithms[j] + ", " + mappings[k])
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('H2 Ground State Energy in different mappings')
pylab.legend(loc='upper right')
pylab.show()

pylab.rcParams['figure.figsize'] = (6, 4)
for k in range(len(mappings)):
    pylab.ylim(-1.14, -1.04)
    pylab.plot(distances, hf_energies, label='Hartree-Fock')
    for j in range(len(algorithms)):
        pylab.plot(distances, energies[k][j], label=algorithms[j])
    pylab.xlabel('Interatomic distance')
    pylab.ylabel('Energy')
    pylab.title('H2 Ground State Energy with {} mapping'.format(mappings[k]))
    pylab.legend(loc='upper right')
    pylab.show()
    
    #pylab.plot(distances, np.subtract(hf_energies, energies[k][1]), label='Hartree-Fock')
    pylab.plot(distances, np.subtract(energies[k][0], energies[k][1]), color=[0.8500, 0.3250, 0.0980], label='VQE')
    pylab.xlabel('Interatomic distance')
    pylab.ylabel('Energy')
    pylab.yscale('log')
    pylab.title('Energy difference from ExactEigensolver with {} mapping'.format(mappings[k]))
    pylab.legend(loc='upper right')
    pylab.show()

