import numpy as np
from qiskit import Aer
from qiskit.transpiler import PassManager

from qiskit_aqua import Operator, QuantumInstance
from qiskit_aqua.algorithms.adaptive import VQE
from qiskit_aqua.algorithms.classical import ExactEigensolver
from qiskit_aqua.components.optimizers import L_BFGS_B
from qiskit_aqua.components.variational_forms import RY

from qiskit_chemistry import FermionicOperator
from qiskit_chemistry.drivers import PySCFDriver, UnitsType

driver = PySCFDriver(atom='H .0 .0 .0; H .0 .0 0.735', unit=UnitsType.ANGSTROM,
                     charge=0, spin=0, basis='sto3g')
molecule = driver.run()

ferOp = FermionicOperator(h1=molecule.one_body_integrals, h2=molecule.two_body_integrals)
qubitOp_jw = ferOp.mapping(map_type='JORDAN_WIGNER', threshold=0.00000001)
qubitOp_jw.chop(10**-10)

# Using exact eigensolver to get the smallest eigenvalue
exact_eigensolver = ExactEigensolver(qubitOp_jw, k=1)
ret = exact_eigensolver.run()

# print(qubitOp_jw.print_operators())

print('The exact ground state energy is: {}'.format(ret['energy']))
print('The Hartree Fock Electron Energy is: {}'.format(molecule.hf_energy - molecule.nuclear_repulsion_energy))

# particle hole transformation
newferOp, energy_shift = ferOp.particle_hole_transformation(num_particles=2)
print('Energy shift is: {}'.format(energy_shift))
newqubitOp_jw = newferOp.mapping(map_type='JORDAN_WIGNER', threshold=0.00000001)
newqubitOp_jw.chop(10**-10)

exact_eigensolver = ExactEigensolver(newqubitOp_jw, k=1)
ret = exact_eigensolver.run()

# print(newqubitOp_jw.print_operators())
print('The exact ground state energy in PH basis is {}'.format(ret['energy']))
print('The exact ground state energy in PH basis is {} (with energy_shift)'.format(ret['energy'] - energy_shift))

# setup VQE 
# setup optimizer, use L_BFGS_B optimizer for example
lbfgs = L_BFGS_B(maxfun=1000, factr=10, iprint=10)

# setup variational form generator (generate trial circuits for VQE)
var_form = RY(newqubitOp_jw.num_qubits, 5, entangler_map = {0: [1], 1:[2], 2:[3]})

# setup VQE with operator, variational form, and optimizer
vqe_algorithm = VQE(newqubitOp_jw, var_form, lbfgs, 'matrix')

backend = Aer.get_backend('statevector_simulator')
quantum_instance = QuantumInstance(backend, pass_manager=PassManager())

results = vqe_algorithm.run(quantum_instance)
print("Minimum value: {}".format(results['eigvals'][0].real))
print("Minimum value: {}".format(results['eigvals'][0].real - energy_shift))
print("Parameters: {}".format(results['opt_params']))

