from lib_q_computer import Qubit, __0, H, CNOT, I, X, Z


def quantum_dense_coding():
    q1 = Qubit(__0)
    q2 = Qubit(__0)

    print("Hadamard on q1")
    q1 = H.on(q1)
    print(q1)
    print()

    print("Bell state")
    bell = CNOT.on([q1, q2])
    print(bell)
    print()

    print("Want to send 00")
    _00 = I.on(bell[0])
    print(_00)
    print()
    print("Want to send 01")
    _01 = X.on(bell[0])
    print(_01)
    print()
    print("Want to send 10")
    _10 = Z.on(bell[0])
    print(_10)
    print()
    print("Want to send 11")
    _11 = X.on(Z.on(bell[0])[0])
    print(_11)
    print()

    print("Bob measures...")
    print("00")
    m00 = H.on(CNOT.on(_00)[0])
    print(m00)
    print()
    print("01")
    m01 = H.on(CNOT.on(_01)[0])
    print(m01)
    print()
    print("10")
    m10 = H.on(CNOT.on(_10)[0])
    print(m10)
    print()
    print("11")
    m11 = H.on(CNOT.on(_11)[0])
    print(m11)


if __name__ == "__main__":
    quantum_dense_coding()