import numpy as np

from qiskit.providers.aer import AerProvider
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from qiskit.circuit import Parameter
from qiskit.tools.visualization import plot_histogram
import qiskit.providers.aer.noise as noise
from qiskit.quantum_info import hellinger_fidelity

qr = QuantumRegister(3)
cr = ClassicalRegister(3)
qc = QuantumCircuit(qr, cr)

tau = Parameter('τ')

qc.h(qr)
qc.cx(qr[0], qr[1])
qc.cx(qr[1], qr[2])

qc.rz(-2 * tau, qr[2])

qc.cx(qr[1], qr[2])
qc.cx(qr[0], qr[1])
qc.h(qr)
qc.draw('mpl')


meas = QuantumCircuit(qr, cr)
meas.measure(qr, cr)
meas.draw('mpl')

qasm_qc = qc.compose(meas)
qasm_qc.draw('mpl')

bind_qasm_qc = qasm_qc.bind_parameters({tau: np.pi/4})

qasm_backend = AerProvider().get_backend('qasm_simulator')
job_qasm = qasm_backend.run(bind_qasm_qc, shots=4096)
counts_qasm = job_qasm.result().get_counts()

print(counts_qasm)
plot_histogram(counts_qasm)

# Error probabilities
prob_1 = 0.001  # 1-qubit gate
prob_2 = 0.01   # 2-qubit gate

# Depolarizing quantum errors
error_1 = noise.depolarizing_error(prob_1, num_qubits=1)
error_2 = noise.depolarizing_error(prob_2, num_qubits=2)

# Add errors to noise model
noise_model = noise.NoiseModel()
noise_model.add_all_qubit_quantum_error(error_1, ['h', 'rz'])
noise_model.add_all_qubit_quantum_error(error_2, ['cx'])

job_qasm_noisy = qasm_backend.run(bind_qasm_qc, shots=4096, noise_model=noise_model)
counts_qasm_noisy = job_qasm_noisy.result().get_counts()
print(counts_qasm_noisy)
plot_histogram(counts_qasm_noisy)

print(hellinger_fidelity(counts_qasm, counts_qasm_noisy))

qc.draw('mpl')

bind_qc = qc.bind_parameters({tau: np.pi/4})
sv_backend = AerProvider().get_backend('statevector_simulator')
job_sv = sv_backend.run(bind_qc)
final_sv = job_sv.result().get_statevector()

np.set_printoptions(suppress=True)
print(final_sv)

unitary_backend = AerProvider().get_backend('unitary_simulator')
job_unitary = unitary_backend.run(bind_qc)
unitary = job_unitary.result().get_unitary()
print(unitary)

for idx, x in np.ndenumerate(unitary):
    if abs(x) > 1e-7:
        print(f"{idx}, {x.real if abs(x.real) > 1e-7 else 0.0} {' + j'+ str(x.imag) if abs(x.imag) > 1e-7 else 0.0}")