# import common packages
import itertools
import logging

import numpy as np

from qiskit import Aer

from qiskit_aqua import Operator, set_aqua_logging, QuantumInstance
from qiskit_aqua.algorithms.adaptive import VQE
from qiskit_aqua.algorithms.classical import ExactEigensolver
from qiskit_aqua.components.optimizers import COBYLA

from qiskit_chemistry.drivers import PySCFDriver, UnitsType
from qiskit_chemistry.core import Hamiltonian, TransformationType, QubitMappingType 
from qiskit_chemistry.aqua_extensions.components.variational_forms import UCCSD
from qiskit_chemistry.aqua_extensions.components.initial_states import HartreeFock

# set_aqua_logging(logging.INFO)

# using driver to get fermionic Hamiltonian
driver = PySCFDriver(atom='Li .0 .0 .0; H .0 .0 1.6', unit=UnitsType.ANGSTROM,
                     charge=0, spin=0, basis='sto3g')
molecule = driver.run()

core = Hamiltonian(transformation=TransformationType.FULL, qubit_mapping=QubitMappingType.PARITY, 
                   two_qubit_reduction=True, freeze_core=True)
algo_input = core.run(molecule)
qubit_op = algo_input.qubit_op

print("Originally requires {} qubits".format(qubit_op.num_qubits))
print(qubit_op)

[symmetries, sq_paulis, cliffords, sq_list] = qubit_op.find_Z2_symmetries()
print('Z2 symmetries found:')
for symm in symmetries:
    print(symm.to_label())
print('single qubit operators found:')
for sq in sq_paulis:
    print(sq.to_label())
print('cliffords found:')
for clifford in cliffords:
    print(clifford.print_operators())
print('single-qubit list: {}'.format(sq_list))

tapered_ops = []
for coeff in itertools.product([1, -1], repeat=len(sq_list)):
    tapered_op = Operator.qubit_tapering(qubit_op, cliffords, sq_list, list(coeff))
    tapered_ops.append((list(coeff), tapered_op))
    print("Number of qubits of tapered qubit operator: {}".format(tapered_op.num_qubits))

ee = ExactEigensolver(qubit_op, k=1)
result = core.process_algorithm_result(ee.run())
for line in result[0]:
    print(line)

smallest_eig_value = 99999999999999
smallest_idx = -1
for idx in range(len(tapered_ops)):
    ee = ExactEigensolver(tapered_ops[idx][1], k=1)
    curr_value = ee.run()['energy']
    if curr_value < smallest_eig_value:
        smallest_eig_value = curr_value
        smallest_idx = idx
    print("Lowest eigenvalue of the {}-th tapered operator (computed part) is {:.12f}".format(idx, curr_value))
    
the_tapered_op = tapered_ops[smallest_idx][1]
the_coeff = tapered_ops[smallest_idx][0]
print("The {}-th tapered operator matches original ground state energy, with corresponding symmetry sector of {}".format(smallest_idx, the_coeff))

# setup initial state
init_state = HartreeFock(num_qubits=the_tapered_op.num_qubits, num_orbitals=core._molecule_info['num_orbitals'],
                    qubit_mapping=core._qubit_mapping, two_qubit_reduction=core._two_qubit_reduction,
                    num_particles=core._molecule_info['num_particles'], sq_list=sq_list)

# setup variationl form
var_form = UCCSD(num_qubits=the_tapered_op.num_qubits, depth=1,
                   num_orbitals=core._molecule_info['num_orbitals'], 
                   num_particles=core._molecule_info['num_particles'],
                   active_occupied=None, active_unoccupied=None, initial_state=init_state,
                   qubit_mapping=core._qubit_mapping, two_qubit_reduction=core._two_qubit_reduction, 
                   num_time_slices=1,
                   cliffords=cliffords, sq_list=sq_list, tapering_values=the_coeff, symmetries=symmetries)

# setup optimizer
optimizer = COBYLA(maxiter=1000)

# set vqe
algo = VQE(the_tapered_op, var_form, optimizer, 'matrix')

# setup backend
backend = Aer.get_backend('statevector_simulator')
quantum_instance = QuantumInstance(backend=backend)

algo_result = algo.run(quantum_instance)

result = core.process_algorithm_result(algo_result)
for line in result[0]:
    print(line)

print("The parameters for UCCSD are:\n{}".format(algo_result['opt_params']))

