import numpy as np
import pylab
from qiskit import LegacySimulators
from qiskit_chemistry import QiskitChemistry
import time

# 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'},
    'algorithm': {'name': ''},
    'initial_state': {'name': 'HartreeFock'},
}
molecule = 'H .0 .0 -{0}; H .0 .0 {0}'
algorithms = [
    {
        'name': 'IQPE',
        'num_iterations': 16,
        'num_time_slices': 3000,
        'expansion_mode': 'trotter',
        'expansion_order': 1,
    },
    {
        'name': 'ExactEigensolver'
    }
]

backends = [
    LegacySimulators.get_backend('qasm_simulator'),
    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)

import concurrent.futures
import multiprocessing as mp
import copy

def subrountine(i, qiskit_chemistry_dict, d, backend, algorithm):
    solver = QiskitChemistry()
    qiskit_chemistry_dict['PYSCF']['atom'] = molecule.format(d/2) 
    qiskit_chemistry_dict['algorithm'] = algorithm
    result = solver.run(qiskit_chemistry_dict, backend=backend)
    return i, d, result['energy'], result['hf_energy']

start_time = time.time()
max_workers = max(4, mp.cpu_count())
with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
    futures = []
    for j in range(len(algorithms)):
        algorithm = algorithms[j]
        backend = backends[j]
        for i in range(steps+1):
            d = start + i*by/steps
            future = executor.submit(
                subrountine, 
                i, 
                copy.deepcopy(qiskit_chemistry_dict), 
                d, 
                backend, 
                algorithm
            )
            futures.append(future)
        for future in concurrent.futures.as_completed(futures):
            i, d, energy, hf_energy = future.result()
            energies[j][i] = energy
            hf_energies[i] = hf_energy
            distances[i] = d
        
print(' --- complete')

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

print("--- %s seconds ---" % (time.time() - start_time))

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

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