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, 'freeze_core': True, 'orbital_reduction': []},
    'algorithm': {'name': 'ExactEigensolver'}
}
molecule = 'H .0 .0 -{0}; Be .0 .0 .0; H .0 .0 {0}'
reductions = [[], [-2, -1], [-3, -2], [-4, -3], [-1], [-2], [-3], [-4]]

pts  = [x * 0.1  for x in range(6, 20)]
pts += [x * 0.25 for x in range(8, 16)]
pts += [4.0]
energies = np.empty([len(reductions), len(pts)])
distances = np.empty(len(pts))

print('Processing step __', end='')
for i, d in enumerate(pts):
    print('\b\b{:2d}'.format(i), end='', flush=True)
    qiskit_chemistry_dict['PYSCF']['atom'] = molecule.format(d) 
    for j in range(len(reductions)):
        qiskit_chemistry_dict['operator']['orbital_reduction'] = reductions[j] 
        solver = QiskitChemistry()
        result = solver.run(qiskit_chemistry_dict)
        energies[j][i] = result['energy']
    distances[i] = d
print(' --- complete')

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


pylab.rcParams['figure.figsize'] = (12, 8)
for j in range(len(reductions)):
    pylab.plot(distances, energies[j], label=reductions[j])
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('BeH2 Ground State Energy')
pylab.legend(loc='upper right')

pylab.rcParams['figure.figsize'] = (12, 8)
for j in range(len(reductions)):
    pylab.plot(distances, np.subtract(energies[j], energies[0]), label=reductions[j])
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('Energy difference compared to no reduction []')
pylab.legend(loc='upper left')

pylab.rcParams['figure.figsize'] = (6, 4)
for j in range(1, len(reductions)):
    pylab.plot(distances, np.subtract(energies[j], energies[0]), color=[1.0, 0.6, 0.2], label=reductions[j])
    pylab.ylim(0, 0.4)
    pylab.xlabel('Interatomic distance')
    pylab.ylabel('Energy')
    pylab.title('Energy difference compared to no reduction []')
    pylab.legend(loc='upper left')
    pylab.show()

e_nofreeze = np.empty(len(pts))
qiskit_chemistry_dict['operator']['orbital_reduction'] = [] 
qiskit_chemistry_dict['operator']['freeze_core'] = False 
for i, d in enumerate(pts):
    print('\b\b{:2d}'.format(i), end='', flush=True)
    qiskit_chemistry_dict['PYSCF']['atom'] = molecule.format(d) 
    solver = QiskitChemistry()
    result = solver.run(qiskit_chemistry_dict)
    e_nofreeze[i] = result['energy']

print(e_nofreeze)

pylab.rcParams['figure.figsize'] = (8, 6)
pylab.plot(distances, energies[0], label='Freeze Core: True')
pylab.plot(distances, e_nofreeze, label='Freeze Core: False')
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('Energy difference, no reduction [], freeze core true/false')
pylab.legend(loc='upper right')
pylab.show()
pylab.title('Energy difference of freeze core True from False')
pylab.plot(distances, np.subtract(energies[0], e_nofreeze), label='Freeze Core: False')
pylab.show()