from qiskit import QuantumCircuit
from qiskit import BasicAer
from qiskit import execute

def encryption(qc, initialString, pk):
    for bit in range(0,len(initialString)):
        if(initialString[bit] == '1'):
            qc.x(bit)
        if(pk[2*bit] == '1'):
            qc.z(bit)
        if(pk[2*bit+1] == '1'):
            qc.x(bit)
    return qc

def decryption(qc, lenCypher, pk):
    for bit in range(0,lenCypher):
        if(pk[2*bit] == '1'):
            qc.z(bit)
        if(pk[2*bit+1] == '1'):
            qc.x(bit)
    return qc

def generate_random_key(length,backend):
    pk = ''
    for i in range(0,2*length):
        qc = QuantumCircuit(1,1)
        qc.h(0)
        qc.measure(0,0)
        counts = execute(qc, backend=backend, shots=1).result().get_counts(qc)
        pk += list(counts.keys())[0]
    return pk

def run_circuit(qc,backend):
    job = execute(qc, backend, shots=100)
    result = job.result()
    return result.get_counts(qc)

initialString = '0110'
pk = generate_random_key(len(initialString),backend)
encrypt = QuantumCircuit(len(initialString),len(initialString))
backend = BasicAer.get_backend('qasm_simulator')

encryption(encrypt,initialString,pk)
encrypt.barrier()
encrypt.measure([0,1,2,3],[0,1,2,3])

print("the encrypted string is: " + str(run_circuit(encrypt,backend)))
encrypt.barrier()

decryption(encrypt,len(encrypt.qubits),pk)

encrypt.barrier()

encrypt.measure([0,1,2,3],[0,1,2,3])

print("the original string is: " + str(run_circuit(encrypt,backend)))

encrypt.draw()