import numpy as np
from scipy.linalg import expm
from qiskit import BasicAer
from qiskit import execute as q_execute
from qiskit import QuantumCircuit, QuantumRegister
from qiskit.quantum_info import state_fidelity
from qiskit.aqua.operators import MatrixOperator, op_converter
from qiskit.aqua.components.initial_states import Custom

num_qubits = 2
evo_time = 1
temp = np.random.random((2 ** num_qubits, 2 ** num_qubits))
h1 = temp + temp.T
qubitOp = MatrixOperator(matrix=h1)
state_in = Custom(num_qubits, state='random')

state_in_vec = state_in.construct_circuit('vector')
groundtruth = expm(-1.j * h1 * evo_time) @ state_in_vec
print('The directly computed groundtruth evolution result state is\n{}.'.format(groundtruth))

groundtruth_evolution = qubitOp.evolve(state_in_vec, evo_time, num_time_slices=0)
print('The groundtruth evolution result as computed by the Dynamics algorithm is\n{}.'.format(groundtruth_evolution))
np.testing.assert_allclose(groundtruth_evolution, groundtruth)

qubit_op = op_converter.to_weighted_pauli_operator(qubitOp)

quantum_registers = QuantumRegister(qubit_op.num_qubits)
circuit = state_in.construct_circuit('circuit', quantum_registers)
circuit += qubit_op.evolve(
    None, evo_time, num_time_slices=1,
    quantum_registers=quantum_registers,
    expansion_mode='suzuki',
    expansion_order=3
)

circuit.draw(output='mpl')

backend = BasicAer.get_backend('statevector_simulator')
job = q_execute(circuit, backend)
circuit_execution_result = np.asarray(job.result().get_statevector(circuit))
print('The evolution result state from executing the Dynamics circuit is\n{}.'.format(circuit_execution_result))

print('Fidelity between the groundtruth and the circuit result states is {}.'.format(
    state_fidelity(groundtruth, circuit_execution_result)
))

