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': 'H .0 .0 -0.3625; H .0 .0 0.3625', 'basis': 'sto3g'},
    'operator': {'name': 'hamiltonian', 'qubit_mapping': 'jordan_wigner',
                 'two_qubit_reduction': False},
    'algorithm': {'name': 'ExactEigensolver'},
    'optimizer': {'name': 'COBYLA', 'maxiter': 10000 },
    'variational_form': {'name': 'RYRZ', 'depth': 3, 'entanglement': 'full'},
    'initial_state': {'name': 'ZERO'}
}
var_forms = ['RYRZ', 'RY']
entanglements = ['full', 'linear']
depths = [x for x in range(3, 11)]

energies = np.empty([len(var_forms), len(entanglements), len(depths)])
hf_energy = None
energy = None
eval_counts = np.empty([len(var_forms), len(entanglements), len(depths)])

solver = QiskitChemistry()
result = solver.run(qiskit_chemistry_dict)
hf_energy = result['hf_energy']
energy = result['energy']
print('Hartree-Fock energy:', hf_energy)
print('FCI energy:', energy)

qiskit_chemistry_dict['algorithm']['name'] = 'VQE' 
print('Processing step __', end='')
for i, d in enumerate(depths):
    print('\b\b{:2d}'.format(i), end='', flush=True)
    qiskit_chemistry_dict['variational_form']['depth'] = d
    for j in range(len(entanglements)):
        qiskit_chemistry_dict['variational_form']['entanglement'] = entanglements[j] 
        for k in range(len(var_forms)):
            qiskit_chemistry_dict['variational_form']['name'] = var_forms[k] 
            solver = QiskitChemistry()
            result = solver.run(qiskit_chemistry_dict)
            energies[k][j][i] = result['energy']
            eval_counts[k][j][i] = result['algorithm_retvals']['eval_count']
print(' --- complete')

print('Depths: ', depths)
print('Energies:', energies)
print('Num evaluations:', eval_counts)

for k in range(len(var_forms)):
    for j in range(len(entanglements)):
        pylab.plot(depths, energies[k][j]-energy, label=var_forms[k]+' + '+entanglements[j])
pylab.xlabel('Variational form depth')
pylab.ylabel('Energy difference')
pylab.yscale('log')
pylab.title('H2 Ground State Energy Difference from Reference')
pylab.legend(loc='upper right')

for k in range(len(var_forms)):
    for j in range(len(entanglements)):
        pylab.plot(depths, eval_counts[k][j], '-o', label=var_forms[k]+' + '+entanglements[j])
pylab.xlabel('Variational form depth')
pylab.ylabel('Evaluations')
pylab.title('VQE number of evaluations')
pylab.legend(loc='upper right')

