from qiskit import *
from qiskit.quantum_info import random_unitary
# Create a random state for qubit C
random_1q_gate = random_unitary(2).to_instruction()

# Create a register for each qubit of interest
qrA = QuantumRegister(1, 'qA')
qrB = QuantumRegister(1, 'qB')
qrC = QuantumRegister(1, 'qC')
anc_top = QuantumRegister(1, 'ancilla_top')
anc_bottom = QuantumRegister(1, 'ancilla_bottom')

# Create a single classical register
cr = ClassicalRegister(2, 'c')
cr_ancilla = ClassicalRegister(1, 'ancilla_bits')
# Create circuit
qc = QuantumCircuit(anc_top, qrC, qrA, qrB, anc_bottom, cr, cr_ancilla)
qc.reset(range(5))
# Start in the |1> state
qc.x(qrC)
# Make C qubit in random quantum state
qc.compose(random_1q_gate, qrC, inplace=True)
# Form Bell pair between A and B
qc.h(qrA)
qc.cx(qrA, qrB)
qc.barrier()

# Bell measurement
qc.cx(qrC, qrA)
qc.h(qrC)
qc.measure([qrC[0],qrA[0]], [0,1])

qc.barrier()
# Corrective Pauli gates
with qc.switch(cr) as case:
    # If returned bit-string is 00 = 0
    with case(0):
        pass
    # If returned bit-string is 01 = 1
    with case(1):
        qc.z(qrB)
    # If returned bit-string is 10 = 2
    with case(2):
        qc.x(qrB)
# If returned bit-string is 11 = 3
    with case(3):
        qc.z(qrB)
        qc.x(qrB)

qc.barrier()
# Apply inverse of unitary that generated state C
qc.compose(random_1q_gate.inverse(), qrB, inplace=True)
# We should be back in the ground state
qc.cx(qrB, anc_bottom)
qc.measure(anc_bottom, cr_ancilla[0])
#qc.measure(qrB, 3)

qc.draw('mpl', fold=-1)