from qiskit.circuit import ClassicalRegister, QuantumRegister, QuantumCircuit, Parameter

theta = Parameter('θ')

qr = QuantumRegister(1, 'q')
qc = QuantumCircuit(qr)
qc.ry(theta, 0)
qc.draw('mpl')

tele_qc = qc.copy()
bell = QuantumRegister(2, 'Bell')
alice = ClassicalRegister(2, 'Alice')
bob = ClassicalRegister(1, 'Bob')
tele_qc.add_register(bell, alice, bob)
tele_qc.draw('mpl')

# create Bell state with other two qubits
tele_qc.barrier()
tele_qc.h(1)
tele_qc.cx(1, 2)
tele_qc.barrier()
tele_qc.draw('mpl')

# alice operates on her qubits
tele_qc.cx(0, 1)
tele_qc.h(0)
tele_qc.barrier()
tele_qc.draw('mpl')

tele_qc.measure([qr[0], bell[0]], alice)
tele_qc.draw('mpl')

full_qc = tele_qc.copy()


#### ANSWER


# Bob's conditional operations
with full_qc.if_test((alice[1], 1)):
    full_qc.x(bell[1])
    
with full_qc.if_test((alice[0], 1)):
    full_qc.z(bell[1])


#### END ANSWER


full_qc.draw('mpl')

full_qc.barrier()
full_qc.measure(bell[1], bob)
full_qc.draw('mpl')

from qiskit_aer.primitives import Sampler
import numpy as np

angle = 5*np.pi/7
sampler = Sampler()
qc.measure_all()
job_static = sampler.run(qc.bind_parameters({theta: angle}))
job_dynamic = sampler.run(full_qc.bind_parameters({theta: angle}))

print(f"Original Dists: {job_static.result().quasi_dists[0].binary_probabilities()}")
print(f"Teleported Dists: {job_dynamic.result().quasi_dists[0].binary_probabilities()}")

from qiskit.result import marginal_counts
from qiskit.visualization import *

# FILL IN CODE HERE

# noticed that the ratio between the teleported states with bob = 0 and bob = 1 is roughly what we expect for the original distribution, so we extracted it and passed in as a dictionary to marginalize counts
bob0 = job_dynamic.result().quasi_dists[0][0] + job_dynamic.result().quasi_dists[0][1] + job_dynamic.result().quasi_dists[0][2] + job_dynamic.result().quasi_dists[0][3]
bob1 = job_dynamic.result().quasi_dists[0][4] + job_dynamic.result().quasi_dists[0][5] + job_dynamic.result().quasi_dists[0][6] + job_dynamic.result().quasi_dists[0][7]
dictionary = {'0':bob0, '1':bob1}
tele_counts = marginal_counts(dictionary)


#### ANSWER

#### END ANSWER



legend = ['Original State', 'Teleported State']
plot_histogram([job_static.result().quasi_dists[0].binary_probabilities(), tele_counts], legend=legend)

import qiskit.tools.jupyter
%qiskit_version_table

