from qiskit import QuantumCircuit, execute, Aer
from qiskit.visualization import plot_histogram
import numpy as np

#prepares the initial quantum states for the BB84 protocol by generating a list of quantum states and their associated bases. 
def prepare_bb84_states(num_qubits):
    states = []
    bases = []
    for _ in range(num_qubits):
        state, basis = generate_bb84_state()
        states.append(state)
        bases.append(basis)
    return states, bases

#generates a single BB84 quantum state and its associated measurement basis.
def generate_bb84_state():
    basis = np.random.randint(2)  # 0 for rectilinear basis, 1 for diagonal basis
    if np.random.rand() < 0.5:
        state = QuantumCircuit(1, 1)
        if basis == 0:
            state.h(0)
        else:
            state.s(0)
            state.h(0)
    else:
        state = QuantumCircuit(1, 1)
        if basis == 0:
            state.x(0)
            state.h(0)
        else:
            state.x(0)
            state.s(0)
            state.h(0)
    return state, basis



#measures the quantum states with the appropriate basis and returns the measurement results
def measure_bb84_states(states, bases):
    measurements = []
    for state, basis in zip(states, bases):
        if basis == 0:
            state.measure(0, 0)
        else:
            state.h(0)
            state.measure(0, 0)
        measurements.append(state)
    return measurements




#sifts the keys by comparing Alice and Bob's measurement bases and bits. It checks if the measurement bases match (the same basis) and appends the corresponding bits to the sifted keys
def sift_keys(alice_bases, bob_bases, alice_bits, bob_bits):
    sifted_alice_bits = []
    sifted_bob_bits = []
    for a_basis, b_basis, a_bit, b_bit in zip(alice_bases, bob_bases, alice_bits, bob_bits):
        if a_basis == b_basis:
            sifted_alice_bits.append(a_bit)
            sifted_bob_bits.append(b_bit)
    return sifted_alice_bits, sifted_bob_bits



#calculates the error rate between Alice's and Bob's sifted bits
def error_rate(alice_bits, bob_bits):
    errors = sum([1 for a, b in zip(alice_bits, bob_bits) if a != b])
    return errors / len(alice_bits)



# simulates the BB84 protocol
num_qubits = 100
num_qubits = 100
alice_states, alice_bases = prepare_bb84_states(num_qubits)
bob_bases = np.random.randint(2, size=num_qubits)
bob_measurements = measure_bb84_states(alice_states, bob_bases)

alice_bits = []
for state in alice_states:
    result = execute(state, Aer.get_backend('qasm_simulator')).result()
    counts = result.get_counts(state)
    max_count = max(counts, key=counts.get)
    alice_bits.append(int(max_count))
bob_bits = []
for state in bob_measurements:
    result = execute(state, Aer.get_backend('qasm_simulator')).result()
    counts = result.get_counts(state)
    max_count = max(counts, key=counts.get)
    bob_bits.append(int(max_count))


sifted_alice_bits, sifted_bob_bits = sift_keys(alice_bases, bob_bases, alice_bits, bob_bits)
error = error_rate(sifted_alice_bits, sifted_bob_bits)
final_key = privacy_amplification(sifted_alice_bits)

print("Sifted key length:", len(sifted_alice_bits))
print("Error rate:", error)
print("Final secret key:", final_key)

