import numpy as np
from qiskit import BasicAer
from qiskit.transpiler import PassManager
from qiskit.aqua import QuantumInstance
from qiskit.aqua.algorithms import VQE, NumPyMinimumEigensolver, IQPE
from qiskit.aqua.operators import WeightedPauliOperator
from qiskit.circuit.library import TwoLocal
from qiskit.aqua.components.optimizers import SPSA
from qiskit.aqua.components.initial_states.var_form_based import VarFormBased

pauli_dict = {
    'paulis': [{"coeff": {"imag": 0.0, "real": -1.052373245772859}, "label": "II"},
              {"coeff": {"imag": 0.0, "real": 0.39793742484318045}, "label": "IZ"},
              {"coeff": {"imag": 0.0, "real": -0.39793742484318045}, "label": "ZI"},
              {"coeff": {"imag": 0.0, "real": -0.01128010425623538}, "label": "ZZ"},
              {"coeff": {"imag": 0.0, "real": 0.18093119978423156}, "label": "XX"}
              ]
}

qubit_op = WeightedPauliOperator.from_dict(pauli_dict)

result_reference = NumPyMinimumEigensolver(qubit_op).run()
print('The reference ground energy level is {}.'.format(result_reference.eigenvalue.real))

random_seed = 0
np.random.seed(random_seed)
backend = BasicAer.get_backend('qasm_simulator')

var_form_depth = 3
var_form = TwoLocal(qubit_op.num_qubits, ['ry', 'rz'], 'cz', reps=var_form_depth)

spsa_max_trials=10
optimizer = SPSA(max_trials=spsa_max_trials)

vqe = VQE(qubit_op, var_form, optimizer)

quantum_instance = QuantumInstance(backend)
result_vqe = vqe.run(quantum_instance)
print('VQE estimated the ground energy to be {}.'.format(result_vqe.eigenvalue.real))

state_in = VarFormBased(var_form, result_vqe.optimal_point)

num_time_slices = 1
num_iterations = 6

iqpe = IQPE(qubit_op, state_in, num_time_slices, num_iterations,
            expansion_mode='suzuki', expansion_order=2,
            shallow_circuit_concat=True)
quantum_instance = QuantumInstance(backend, shots=100, seed_simulator=random_seed, seed_transpiler=random_seed)
result_iqpe = iqpe.run(quantum_instance)
print("Continuing with VQE's result, IQPE estimated the ground energy to be {}.".format(
    result_iqpe.eigenvalue.real))