# useful additional packages 
import matplotlib.pyplot as plt
import copy
%matplotlib inline
import numpy as np
from qiskit import Aer
from qiskit.aqua import QuantumInstance
from qiskit.aqua.algorithms import NumPyMinimumEigensolver, VQE
from qiskit.circuit.library import TwoLocal
from qiskit.aqua.components.optimizers import SPSA
from qiskit.chemistry.core import Hamiltonian, TransformationType, QubitMappingType
from qiskit.chemistry.drivers import HDF5Driver
import warnings
warnings.filterwarnings('ignore')

# setup qiskit.chemistry logging
import logging
from qiskit.chemistry import set_qiskit_chemistry_logging
set_qiskit_chemistry_logging(logging.ERROR) # choose among DEBUG, INFO, WARNING, ERROR, CRITICAL and NOTSET

# from qiskit import IBMQ
# provider = IBMQ.load_account()

# First, we use classical eigendecomposition to get ground state energy (including nuclear repulsion energy) as reference.

driver = HDF5Driver('H2/H2_equilibrium_0.735_sto-3g.hdf5')
molecule = driver.run()
operator =  Hamiltonian(transformation=TransformationType.FULL,
                 qubit_mapping=QubitMappingType.PARITY,
                 two_qubit_reduction=True,
                 freeze_core=False,
                 orbital_reduction=None)
qubit_op, aux_ops = operator.run(molecule)
result = NumPyMinimumEigensolver(qubit_op).run()
print('Ground state energy (classical): {:.12f}'.format(result.eigenvalue.real))

# Second, we use variational quantum eigensolver (VQE)

var_form = TwoLocal(qubit_op.num_qubits, ['ry', 'rz'], 'cz', reps=3, entanglement='full')
optimizer = SPSA(max_trials=350)
algo = VQE(qubit_op, var_form, optimizer, max_evals_grouped=1)
result = algo.run(QuantumInstance(Aer.get_backend('statevector_simulator')))
result = operator.process_algorithm_result(result)
        
print('Ground state energy (quantum)  : {:.12f}'.format(result.energy))
print("====================================================")
# You can also print out other info
print(result)

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

if molecule == 'LiH':
    mol_distances = np.arange(0.6, 5.1, 0.1)
else:
    mol_distances = np.arange(0.2, 4.1, 0.1)

algos = ['NumPyMinimumEigensolver', 'VQE']
energy = np.zeros((len(algos), len(mol_distances)))
for j, algo_name in enumerate(algos):
    print("Using {}".format(algo_name))
    for i, dis in enumerate(mol_distances):
        print("Processing atomic distance: {:1.1f} Angstrom".format(dis), end='\r') 
        if molecule == 'LiH':
            operator =  Hamiltonian(transformation=TransformationType.FULL,
                             qubit_mapping=QubitMappingType.PARITY,
                             two_qubit_reduction=True,
                             freeze_core=True,
                             orbital_reduction=[-3, -2])
        else:
            operator =  Hamiltonian(transformation=TransformationType.FULL,
                             qubit_mapping=QubitMappingType.PARITY,
                             two_qubit_reduction=True,
                             freeze_core=False,
                             orbital_reduction=None)
            
        driver = HDF5Driver("{}/{:1.1f}_sto-3g.hdf5".format(molecule, dis))
        qubit_op, aux_ops = operator.run(driver.run())
        if algo_name == 'VQE':
            if molecule == 'LiH':
                var_form = TwoLocal(qubit_op.num_qubits, ['ry', 'rz'], 'cz', reps=5, entanglement='full')
                optimizer = SPSA(max_trials=2500)
            else:
                var_form = TwoLocal(qubit_op.num_qubits, ['ry', 'rz'], 'cz', reps=3, entanglement='full')
                optimizer = SPSA(max_trials=350)
            
            algo = VQE(qubit_op, var_form, optimizer, max_evals_grouped=1)
            result = algo.run(QuantumInstance(Aer.get_backend('qasm_simulator'), shots=1024))
        else:
            result = NumPyMinimumEigensolver(qubit_op).run()
        
        result = operator.process_algorithm_result(result)
        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()

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright