import warnings
import numpy as np 
from numpy import pi
from matplotlib import MatplotlibDeprecationWarning
from qiskit import QuantumRegister, QuantumCircuit, execute, BasicAer
from qiskit_optimization import QuadraticProgram
from qiskit.visualization import plot_state_qsphere
from qiskit.quantum_info.operators import Operator

qreg_q = QuantumRegister(3, 'q')
qc_init = QuantumCircuit(qreg_q)
qc_init.h(qreg_q[0])
qc_init.h(qreg_q[1])
qc_init.h(qreg_q[2])
qc_init.barrier()
psi_0 = execute(qc_init, BasicAer.get_backend('statevector_simulator')).result().get_statevector(qc_init, \
                                              decimals=3)
superpos_state = np.ones(8, dtype=complex) / np.sqrt(8)
print("Superposition state: ",np.allclose(psi_0, superpos_state, rtol=0.01))
warnings.filterwarnings("ignore", category=MatplotlibDeprecationWarning)
plot_state_qsphere(psi_0)

def pretty_print(x):
    """ print with 1 decimal only """
    with np.printoptions(precision=1, suppress=True):
        print(x)

qc_evol = QuantumCircuit(qreg_q)
qc_evol.cx(qreg_q[1], qreg_q[0])
qc_evol.rz(-0.06942544294642758, qreg_q[0])
qc_evol.cx(qreg_q[1], qreg_q[0])
qc_evol.rz(-0.20827632883928274, qreg_q[0])
qc_evol.cx(qreg_q[2], qreg_q[1])
qc_evol.rz(1.2496579730356965, qreg_q[1])
qc_evol.cx(qreg_q[2], qreg_q[1])
qc_evol.rz(-0.20827632883928274, qreg_q[1])
qc_evol.barrier()
qc_evol.rx(0.7530122052855767, qreg_q[0])
qc_evol.rx(0.7530122052855767, qreg_q[1])
qc_evol.rx(0.7530122052855767, qreg_q[2])
U = execute(qc_evol, BasicAer.get_backend('unitary_simulator')).result().get_unitary(qc_evol, decimals=3)
pretty_print(U)

op = Operator(qc_evol) # to make it easy to multiply
psi_g = op.data.dot(psi_0)
pretty_print(psi_g)
warnings.filterwarnings("ignore", category=MatplotlibDeprecationWarning)
plot_state_qsphere(psi_g, figsize=(10,10))

qc_qaoa = qc_init.compose(qc_evol)
qc_qaoa.measure_all()
qc_qaoa.draw(output='latex') # ref to fig:circuit_ideal