import numpy as np
import pylab
from qiskit_chemistry import QiskitChemistry

# Input dictionary to configure Qiskit Chemistry for the chemistry problem.
# Note: In order to allow this to run reasonably quickly it takes advantage
#       of the ability to freeze core orbitals and remove unoccupied virtual
#       orbitals to reduce the size of the problem. The result without this
#       will be more accurate but it takes rather longer to run.

# qiskit_chemistry_dict_eigen uses classical approach to produce the reference ground state energy.
qiskit_chemistry_dict_eigen = {
    'driver': {'name': 'PYSCF'},
    'PYSCF': {'atom': '', 'basis': 'sto3g'},
    'algorithm': {'name': 'ExactEigensolver'},
    'operator': {'name':'hamiltonian','freeze_core': True, 'orbital_reduction': [-3, -2], 'qubit_mapping': 'parity', 'two_qubit_reduction': True},
}

# qiskit_chemistry_dict_vqe uses quantum approach to evaluate the ground state energy.
qiskit_chemistry_dict_vqe = {
    'driver': {'name': 'PYSCF'},
    'PYSCF': {'atom': '', 'basis': 'sto3g'},
    'algorithm': {'name': 'VQE', 'operator_mode': 'matrix'},
    'operator': {'name':'hamiltonian','freeze_core': True, 'orbital_reduction': [-3, -2], 'qubit_mapping': 'parity', 'two_qubit_reduction': True},
    'optimizer': {'name': 'COBYLA', 'maxiter': 20000},
    'variational_form': {'name': 'RYRZ', 'depth': 10},
    'backend': {'name': 'statevector_simulator'}
}

# tested molecular, LiH
molecule = 'Li .0 .0 -{0}; H .0 .0 {0}'

# choose one of configurations above for experiments
qiskit_chemistry_dict = qiskit_chemistry_dict_eigen

# configure distance between two atoms
pts  = [x * 0.1  for x in range(6, 20)]
pts += [x * 0.25 for x in range(8, 16)]
pts += [4.0]
distances   = np.empty(len(pts))
hf_energies = np.empty(len(pts))
energies    = np.empty(len(pts))
dipoles     = 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/2) 
    solver = QiskitChemistry()
    result = solver.run(qiskit_chemistry_dict)
    distances[i] = d
    hf_energies[i] = result['hf_energy']
    energies[i] = result['energy']
    dipoles[i]  = result['total_dipole_moment'] / 0.393430307
print(' --- complete')

print('Distances: ', distances)
print('HF Energies:', hf_energies)
print('Energies:', energies)
print('Dipole moments:', dipoles)

pylab.plot(distances, hf_energies, label='HF')
pylab.plot(distances, energies, label='Computed')
pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('LiH Ground State Energy')
pylab.legend(loc='upper right')

pylab.plot(distances, dipoles)
pylab.xlabel('Interatomic distance')
pylab.ylabel('Moment debye')
pylab.title('LiH Dipole Moment')

