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

q = QuantumRegister(2)
circ = QuantumCircuit(q)

circ.h(q[0])
circ.cx(q[0], q[1])

# print(circ)

circ.measure_all()

# print(circ)

# a and b are the qubits in circ
def bell(a, b, circ):
    circ.h(a)
    circ.cx(a, b)
    return circ

def bell_inv(a, b, circ):
    circ.cx(a, b)
    circ.h(a)
    return circ

def tele():
    q = QuantumRegister(3)
    b = ClassicalRegister(3)
    circ = QuantumCircuit(q, b)

    
    circ = bell(q[1], q[2], circ)

    # circ.x(q[0])
    circ.h(q[0])
        
    circ = bell_inv(q[0], q[1], circ)

    circ.measure(q[0], b[0])
    circ.measure(q[1], b[1])
    with circ.if_test ((b[1], 1)):
        circ.x(q[2])
    
    with circ.if_test ((b[0], 1)):
        circ.z(q[2])

    circ.h(q[2])
    circ.measure(q[2], b[2])
    return circ


circ1 = tele() 

print(circ1)



from qiskit_aer import AerSimulator
from qiskit.primitives import BackendSamplerV2

# # circ1 = tele()
aer = AerSimulator()
sampler = BackendSamplerV2(backend=aer)

result = sampler.run([circ1], shots=1024).result()

print(result)

print(result[0].data)
print(result[0].data.keys())

# # use this if we use measure_all method

print(result[0].data.c0.get_counts())


# # for reg_name in result[0].data.keys():
# #     counts = result[0].data[reg_name].get_counts()
# #     print(reg_name, counts)