from random import randint,sample
def get_random_bitstring(sz):
    return [randint(0,1) for i in range(sz)]

from qiskit import (QuantumCircuit, execute, Aer)

default_backend = Aer.get_backend('qasm_simulator')

X_BASIS = 0
Z_BASIS = 1

def encode_key_bit(qbit, tx_bit, tx_base):
    '''
    Append encoding of the specified classical bit 'tx_bit' into the specified
    base 'tx_base' to the quantum circuit
    '''
    
    # flip the quantum 0 to quantum 1 if we received a '1' bit; otherwise leave 
    # it as a quantum 0.
    if tx_bit:
        qbit.x(0)
    
    # if we're transmitting using the x-basis, then leave it in the x-basis;
    # otherwise flip the qubit into the z-basis    
    if tx_base == Z_BASIS:
        qbit.h(0)            
        
    return qbit

def decode_key_bit(qbit, rx_base):
    '''
    Append measurement of the specified qubit 'qbit' in the specified base 'rx_base'
    to the quantum circuit
    '''

    # flip the qubit if we're measureing in the Z_BASIS so if the qubit
    # is originally in the X_BASIS, it will not be in superposition. Then
    # go ahead and measure it.
    if rx_base == Z_BASIS:
        qbit.h(0)    
    qbit.measure([0],[0])
    return qbit

datum = 1
tx_base = 1
rx_base = 1

test_qc = QuantumCircuit(1,1)
test_qc = encode_key_bit(test_qc,datum,tx_base)
test_qc.barrier(0)
test_qc = decode_key_bit(test_qc,rx_base)

test_qc.draw('mpl')

num_bits = 32

# set up Alice's transmission bits
random_tx_bits = get_random_bitstring(num_bits)
random_tx_bases = get_random_bitstring(num_bits)

# set up Bob's receiving bits
random_rx_bases = get_random_bitstring(num_bits)
decoded_rx_bits = []

class SharedChannel:
    
    def __init__(self):
        self.datum = None
        self.active = False
    
    def transmit_datum(self, datum):
        self.datum = datum
        self.active = True
        
    def receive_datum(self):
        return_datum = self.datum
        self.active = False
        self.datum = []
        return return_datum

class SharedQuantumChannel(SharedChannel):
    def __init__(self):
        super().__init__()

class SharedClassicalChannel(SharedChannel):
    def __init__(self):
        super().__init__()
        
    def transmit_data(self, datum):
        super().transmit_datum(datum)
        
    def receive_data(self):
        return super().receive_datum()

# create shared quantum and classical channels that Alice and Bob will
# use to communicate to each other
sqc = SharedQuantumChannel()
scc = SharedClassicalChannel()

scc.transmit_datum('1')
print(scc.receive_datum())

# Clear decoded bits list and transmit the qbits one by one until the 
# receiver has measured every one    

decoded_rx_bits = []
for i in range(num_bits):
    
    # create a qubit, encode it, transmit it, receive it, 
    # and decode it
    qbit = QuantumCircuit(1,1)
    tx_qbit = encode_key_bit(qbit, random_tx_bits[i], random_tx_bases[i])
    sqc.transmit_datum(tx_qbit)
    rx_qbit = sqc.receive_datum()
    full_circuit = decode_key_bit(rx_qbit, random_rx_bases[i])

    job = execute(full_circuit, backend=default_backend, shots=1)
    results = job.result()
    counts = results.get_counts()    
    decoded_rx_bits.append(int(max(counts, key=counts.get)))
    
# decoder reports rx bases to transmitter
scc.transmit_data(random_rx_bases)
rx_bases_used = scc.receive_data()

# sender determines which random rx bases were right
assert(len(rx_bases_used) == len(random_tx_bits))
rx_base_hits = [rx_bases_used[i] == random_tx_bases[i] for i in range(len(random_tx_bits))]

# sender notifies receiver of key hits
scc.transmit_data(rx_base_hits)
rcvd_rx_base_hits = scc.receive_data()

# sender and receiver drop non-hit key bits
sender_key_candidate   = [j for i,j in enumerate(random_tx_bits) if rx_base_hits[i]]
receiver_key_candidate = [j for i,j in enumerate(decoded_rx_bits) if rcvd_rx_base_hits[i]]
assert(sender_key_candidate == receiver_key_candidate)

# the receiver chooses some verification bits (1/4th of the decoded key bits)
num_verify_bits = len(receiver_key_candidate)//4
verify_bits_ix = sample(range(len(receiver_key_candidate)), num_verify_bits)
verify_bits = {i:receiver_key_candidate[i] for i in verify_bits_ix}

# the receiver shares verification bits with sender
scc.transmit_data(verify_bits)
chk_verify_bits = scc.receive_data()
assert(verify_bits == chk_verify_bits)

# verify the check bits
verified = all([verify_bits[i] == sender_key_candidate[i] for i in verify_bits.keys()])

# if the check bits check out, strip them out of the key and we're done
if verified:
    sender_key   = [j for i,j in enumerate(sender_key_candidate) if i not in chk_verify_bits.keys()]
    receiver_key = [j for i,j in enumerate(receiver_key_candidate) if i not in verify_bits.keys()]
    assert(sender_key == receiver_key)
    
print('skey: ' + str(sender_key))
print('rkey: ' + str(receiver_key))

