import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit import Aer, execute
from qiskit.circuit import Gate

# Define no. of message qubits
n = 10

# Alice's end
alice_prepare = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
alice_hadamard = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
circ = QuantumCircuit(n, n)

# Preparation
circ.reset(range(n))
for i in range(n):
    if alice_prepare[i] == 1:
        circ.x(i)
circ.barrier()

# Hadamard Transform
for i in range(n):
    if alice_hadamard[i] == 1:
        circ.h(i)
circ.barrier()

# Eavesdropper
circ.append(Gate(name="$Eve$", num_qubits=n, params=[]), range(n))

# Bob's end
bob_hadamard = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
circ.barrier()

# Hadamard Transform
for i in range(n):
    if bob_hadamard[i] == 1:
        circ.h(i)
circ.barrier()

# Measurement
circ.measure(range(n), range(n))
circ.draw('mpl', reverse_bits=True, scale=0.45)

# Alice's end
alice_prepare = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
alice_hadamard = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
circ = QuantumCircuit(n, n)

# Preparation
circ.reset(range(n))
for i in range(n):
    if alice_prepare[i] == 1:
        circ.x(i)
circ.barrier()

# Hadamard Transform
for i in range(n):
    if alice_hadamard[i] == 1:
        circ.h(i)
circ.barrier()

# Bob's end
bob_hadamard = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])

# Hadamard Transform
for i in range(n):
    if bob_hadamard[i] == 1:
        circ.h(i)
circ.barrier()

# Measurement
circ.measure(range(n), range(n))
circ.draw('mpl', reverse_bits=True, scale=0.45)

# Execute Circuit
backend = Aer.get_backend('qasm_simulator')
job = execute(circ, backend, shots=1)
counts = job.result().get_counts(circ)
bob_measure = list(counts.keys())[0][::-1]

# Print Data
print("".join(["Alice Prepared:          ","".join([str(i) for i in alice_prepare])]))
print("".join(["Alice H Transforms:      ","".join([str(i) for i in alice_hadamard])]))
print("".join(["Bob H Transforms:        ","".join([str(i) for i in bob_hadamard])]))
print("".join(["Bob Measurements:        ", bob_measure]))

# Extract Key
alice_key = []
bob_key = []
for i in range(n):
    if alice_hadamard[i]==bob_hadamard[i]:
        alice_key.append(int(alice_prepare[i]))
        bob_key.append(int(bob_measure[i]))
        
# Print Keys
print("".join(["Alice extracts key as:   ","".join([str(i) for i in alice_key])]))
print("".join(["Bob extracts key as:     ","".join([str(i) for i in bob_key])]))

# Alice's end
alice_prepare = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
alice_hadamard = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
circ = QuantumCircuit(n, 2*n)

# Preparation
circ.reset(range(n))
for i in range(n):
    if alice_prepare[i] == 1:
        circ.x(i)
circ.barrier()

# Hadamard Transform
for i in range(n):
    if alice_hadamard[i] == 1:
        circ.h(i)
circ.barrier()

# Eavesdropper
circ.measure(range(n), range(n))

# Bob's end
bob_hadamard = np.random.choice([0,1], size=(n,), p=[0.5, 0.5])
circ.barrier()

# Hadamard Transform
for i in range(n):
    if bob_hadamard[i] == 1:
        circ.h(i)
circ.barrier()

# Measurement
circ.measure(range(n), range(n,2*n))
circ.draw('mpl', reverse_bits=True, scale=0.45)

# Execute Circuit
backend = Aer.get_backend('qasm_simulator')
job = execute(circ, backend, shots=1)
counts = job.result().get_counts(circ)
measure_data = list(counts.keys())[0][::-1]
eve_measure = measure_data[0:n]
bob_measure = measure_data[n:2*n]

# Print Data
print("".join(["Alice Prepared:          ","".join([str(i) for i in alice_prepare])]))
print("".join(["Alice H Transforms:      ","".join([str(i) for i in alice_hadamard])]))
print("".join(["Eve Measurements:        ", eve_measure]))
print("".join(["Bob H Transforms:        ","".join([str(i) for i in bob_hadamard])]))
print("".join(["Bob Measurements:        ", bob_measure]))

# Extract Key
alice_key = []
bob_key = []
for i in range(n):
    if alice_hadamard[i]==bob_hadamard[i]:
        alice_key.append(int(alice_prepare[i]))
        bob_key.append(int(bob_measure[i]))
        
# Print Keys
print("".join(["Alice extracts key as:   ","".join([str(i) for i in alice_key])]))
print("".join(["Bob extracts key as:     ","".join([str(i) for i in bob_key])]))