# useful additional packages 
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np

from qiskit_acqua_chemistry import ACQUAChemistry
import warnings
warnings.filterwarnings('ignore')

import logging
logger = logging.getLogger()
# logger.setLevel(logging.DEBUG) # uncomment it to see detailed logging

################# import Qconfig and set APIToken and API url and  prepare backends ############
try:
    import sys
    sys.path.append("../../") # go to parent dir
    import Qconfig
except Exception as e:
    print(e)
    
from qiskit import register, available_backends

#set api
APItoken=getattr(Qconfig, 'APItoken', None)
url = Qconfig.config.get('url', None)
hub = Qconfig.config.get('hub', None)
group = Qconfig.config.get('group', None)
project = Qconfig.config.get('project', None)
try:
    register(APItoken, url, hub, group, project)
except Exception as e:
    print(e)

print("Backends: {}".format(available_backends()))

# First, we use classical eigendecomposition to get ground state energy (including nuclear repulsion energy) as reference.
acqua_chemistry_dict = {
    'driver': {'name': 'HDF5'},
    'HDF5': {'hdf5_input': 'H2/H2_equilibrium_0.735_sto-3g.hdf5'},
    'operator': {'name':'hamiltonian', 
                 'qubit_mapping': 'parity', 
                 'two_qubit_reduction': True},
    'algorithm': {'name': 'ExactEigensolver'}
}
solver = ACQUAChemistry()
result = solver.run(acqua_chemistry_dict)
print('Ground state energy (classical): {:.12f}'.format(result['energy']))

# Second, we use variational quantum eigensolver (VQE)
acqua_chemistry_dict['algorithm']['name'] = 'VQE'
acqua_chemistry_dict['optimizer'] = {'name': 'SPSA', 'max_trials': 350}
acqua_chemistry_dict['variational_form'] = {'name': 'RYRZ', 'depth': 3, 'entanglement':'full'}
acqua_chemistry_dict['backend'] = {'name': 'local_statevector_simulator'}

solver = ACQUAChemistry()
result = solver.run(acqua_chemistry_dict)
print('Ground state energy (quantum)  : {:.12f}'.format(result['energy']))
print("====================================================")
# You can also print out other info in the field 'printable'
for line in result['printable']:
    print(line)

# select H2 or LiH to experiment with
molecule='H2'

acqua_chemistry_dict = {
    'driver': {'name': 'HDF5'},
    'HDF5': {'hdf5_input': ''},
    'operator': {'name':'hamiltonian', 
                 'qubit_mapping': 'parity', 
                 'two_qubit_reduction': True},
    'algorithm': {'name': ''},
    'optimizer': {'name': 'SPSA', 'max_trials': 350},
    'variational_form': {'name': 'RYRZ', 'depth': 3, 'entanglement':'full'}
}

# choose which backend want to use
# backend = {'name': 'local_statevector_simulator'}
backend = {'name': 'local_qasm_simulator', 'shots': 1024}
algos = ['ExactEigensolver', 'VQE']

if molecule == 'LiH':
    mol_distances = np.arange(0.6, 5.1, 0.1)
    acqua_chemistry_dict['operator']['freeze_core'] = True
    acqua_chemistry_dict['operator']['orbital_reduction'] = [-3, -2]
    acqua_chemistry_dict['optimizer']['max_trials'] = 2500
    acqua_chemistry_dict['variational_form']['depth'] = 5
else:
    mol_distances = np.arange(0.2, 4.1, 0.1)

energy = np.zeros((len(algos), len(mol_distances)))

for j, algo in enumerate(algos):
    acqua_chemistry_dict['algorithm']['name'] = algo
    if algo == 'ExactEigensolver':
        acqua_chemistry_dict.pop('backend', None)
    elif algo == 'VQE':
        acqua_chemistry_dict['backend'] = backend
    print("Using {}".format(algo))
    for i, dis in enumerate(mol_distances):
        print("Processing atomic distance: {:1.1f} Angstrom".format(dis), end='\r')
        acqua_chemistry_dict['HDF5']['hdf5_input'] = "{}/{:1.1f}_sto-3g.hdf5".format(molecule, dis)
        result = solver.run(acqua_chemistry_dict)
        energy[j][i] = result['energy']
    print("\n")

for i, algo in enumerate(algos):
    plt.plot(mol_distances, energy[i], label=algo)
plt.xlabel('Atomic distance (Angstrom)')
plt.ylabel('Energy')
plt.legend()
plt.show()