#!pip install qiskit

from qiskit import QuantumCircuit, Aer, transpile, assemble, execute, IBMQ
from qiskit.visualization import plot_histogram

import numpy as np
import random
import math

import warnings
warnings.filterwarnings('ignore')


# Determine the amount of entanglement between these bits using the CHSH value
def entanglement_amount(alice_choices, alice_bits, bob_choices, bob_bits):
    
  # count the different measurement results
  # rows correspond to Alice and Bob's circuit choices: 00, 02, 20, 22
  # NOTE: We do not consider circuits 1 or 3 for this test
  # columns correspond to Alice and Bob's qubit measurements: 00, 01, 10, and 11
  circuits = {'00': 0, '02': 1, '20': 2, '22': 3}
  counts = [[0]*4 for i in range(4)]
  for i in range(len(alice_choices)):
    circuit = str(alice_choices[i]) + str(bob_choices[i])
    state = int(alice_bits[i]) + 2*int(bob_bits[i])

    if circuit in circuits: counts[circuits[circuit]][state] += 1


  # expectation values calculated by 
  # adding times Alice and Bob's bits agreed and
  # subtracting times Alice and Bob's bits disagreed
  expectations = []

  for circuit in range(4):
    expectations += [counts[circuit][0] + counts[circuit][3] - counts[circuit][1] - counts[circuit][2]]
    total = sum(counts[circuit])

    if total != 0: expectations[circuit] /= total


  # returning CHSH correlation
  return expectations[0] - expectations[1] + expectations[2] + expectations[3]
  

print("Libraries Imported Successfully!")

alice_bob_qubits = QuantumCircuit(2, 2)

alice_bob_qubits.# COMPLETE THIS CODE
alice_bob_qubits.# COMPLETE THIS CODE
alice_bob_qubits.# COMPLETE THIS CODE
alice_bob_qubits.# COMPLETE THIS CODE

alice_bob_qubits.draw()

alice_option_1 = QuantumCircuit(1, 1)

alice_option_1.h(0)

alice_option_1.measure(0, 0)

alice_option_2 = QuantumCircuit(1, 1)

alice_option_2.rz(math.pi# COMPLETE THIS LINE

alice_option_2.h(0)

alice_option_2.measure(0, 0)

alice_option_3 = QuantumCircuit(1, 1)

alice_option_3.# COMPLETE THIS LINE

alice_alice_option_3.h(0)

alice_option_3.measure(0, 0)

alice_options = [alice_option_1, # COMPLETE THIS LINE WITH THE OTHER OPTIONS
alice_choice = random.randint(0, 2)

alice_circuit = options[alice_choice]
alice_circuit.draw()

alice_bob_qubits = alice_bob_qubits.compose(alice_circuit, qubits = 0, clbits = 0)

alice_bob_qubits.draw()

bob_option_1 = QuantumCircuit(1, 1)

bob_option_1.# COMPLETE THIS LINE
bob_option_1.# COMPLETE THIS LINE

bob_option_1.measure(0, 0)

bob_option_2 # COMPLETE THIS CODE

bob_option_3 # COMPLETE THIS CODE

bob_options = # COMPLETE THIS LINE
bob_choice = random.randint(0, 2)

bob_circuit = options[bob_choice]
bob_circuit.draw()

alice_bob_qubits = alice_bob_qubits.compose(bob_circuit, qubits = 1, clbits = 1)

alice_bob_qubits.draw()

backend = Aer.get_backend('qasm_simulator')
job = execute(alice_bob_qubits, backend = backend, shots = 1024)


result = job.result()
counts = result.get_counts()

bits = random.choices(list(counts.keys()), weights = counts.values(), k = 1)[0]
alice_bits = bits[0]
bob_bits = bits[1]


plot_histogram(counts)

alice_key = []
alice_mismatched_choices = []
alice_mismatched_choice_bits = []

bob_key = []
bob_mismatched_choices = []
bob_mismatched_choice_bits = []


# MATCHING CHOICE
if alice_options[alice_choice] == bob_options[bob_choice]:
  alice_key += [int(alice_bits[0])]
  bob_key += [1 - int(bob_bits[0])]

# MISMATCHING CHOICE
else:
  alice_mismatched_choices += [alice_choice]
  bob_mismatched_choices += [bob_choice]

  alice_mismatched_choice_bits += [alice_bits[0]]
  bob_mismatched_choice_bits += [bob_bits[0]]


entanglement = entanglement_amount(alice_mismatched_choices, alice_mismatched_choice_bits, bob_mismatched_choices, bob_mismatched_choice_bits)
print("Entanglement of Mismatched Choices: " + str(entanglement))
print("Alice's Key: " + str(alice_key))
print("Bob's Key: " + str(bob_key))

#========
# STEP 1
#========
alice_bob_qubits = QuantumCircuit(2, 2)

# COMPLETE THIS CODE


#========
# STEP 2
#========
alice_choice = random.randint(0, 2)
alice_circuit = # COMPLETE THIS LINE


#========
# STEP 3
#========
alice_bob_qubits = alice_bob_qubits.compose(# COMPLETE THIS LINE


#========
# STEP 4
#========
bob_choice = random.randint(0, 2)
bob_circuit = # COMPLETE THIS LINE


#========
# STEP 5
#========
alice_bob_qubits = alice_bob_qubits.compose(# COMPLETE THIS LINE


#======================
# SIMULATE THE CIRCUIT
#======================
backend = Aer.get_backend('qasm_simulator')
job = execute(alice_bob_qubits, backend = backend, shots = 1024)
result = job.result()
counts = result.get_counts()

bits = random.choices(list(counts.keys()), weights = counts.values(), k = 1)[0]
alice_bits = bits[0]
bob_bits = bits[1]
plot_histogram(counts)


#========
# STEP 6
#========
alice_key = []
alice_mismatched_choices = []
alice_mismatched_choice_bits = []

bob_key = []
bob_mismatched_choices = []
bob_mismatched_choice_bits = []


# MATCHING CHOICE
if alice_options[alice_choice] == bob_options[bob_choice]:
  alice_key += [int(# COMPLETE THIS LINE
  bob_key += [1 - int(# COMPLETE THIS LINE

# MISMATCHING CHOICE
else:
  alice_mismatched_choices += [# COMPLETE THIS LINE
  bob_mismatched_choices += [# COMPLETE THIS LINE

  alice_mismatched_choice_bits += [# COMPLETE THIS LINE
  bob_mismatched_choice_bits += [# COMPLETE THIS LINE


entanglement = entanglement_amount(alice_mismatched_choices, alice_mismatched_choice_bits, bob_mismatched_choices, bob_mismatched_choice_bits)
print("Entanglement of Mismatched Choices: " + str(entanglement))
print("Alice's Key: " + str(alice_key))
print("Bob's Key: " + str(bob_key))

n = 100

alice_bob_qubits = []

for i in range(# COMPLETE THIS LINE
  alice_bob_qubits += [QuantumCircuit(# COMPLETE THIS LINE

  alice_bob_qubits[i].# COMPLETE THIS LINE
  alice_bob_qubits[i].# COMPLETE THIS LINE
  alice_bob_qubits[i].# COMPLETE THIS LINE
  alice_bob_qubits[i].# COMPLETE THIS LINE

alice_choices = []
alice_circuits = []

for # COMPLETE THIS LINE
  alice_choices += [random.randint(0, 2)]
  alice_circuits += [# COMPLETE THIS LINE

# COMPLETE THIS LINE
  alice_bob_qubits[i] = alice_bob_qubits[i].compose(# COMPLETE THIS LINE

# COMPLETE THIS CODE

# COMPLETE THIS CODE

backend = Aer.get_backend('qasm_simulator')
job = execute(alice_bob_qubits, backend = backend, shots = 1)
result = job.result()
counts = result.get_counts()

alice_bits = []
bob_bits = []
for i in range(n):
  bits = list(counts[i].keys())[0]
  alice_bits += [bits[0]]
  bob_bits += [bits[1]]

alice_key = []
alice_mismatched_choices = []
alice_mismatched_choice_bits = []

bob_key = []
bob_mismatched_choices = []
bob_mismatched_choice_bits = []

for # COMPLETE THIS LINE

  # MATCHING CHOICE
  if alice_options[alice_choices[i]] == # COMPLETE THIS LINE
    alice_key += [int(alice_bits[i])]
    bob_key += [1 - int(# COMPLETE THIS LINE

  # MISMATCHING CHOICE
  else:
    alice_mismatched_choices += [alice_choices[i]]
    bob_mismatched_choices += [ # COMPLETE THIS LINE

    alice_mismatched_choice_bits += [ # COMPLETE THIS LINE
    bob_mismatched_choice_bits +=  [# COMPLETE THIS LINE


entanglement = entanglement_amount(alice_mismatched_choices, alice_mismatched_choice_bits, bob_mismatched_choices, bob_mismatched_choice_bits)
print("Entanglement of Mismatched Choices: " + str(entanglement))
print("Alice's Key: " + str(alice_key))
print("Bob's Key: " + str(bob_key))
print("Key Length: " + str(len(bob_key)))
print("Number of Disagreeing Key Bits: " + str(sum([alice_key[i] != bob_key[i] for i in range(len(alice_key))])))

#========
# STEP 1
#========
n = 100


#========
# STEP 2
#========



#========
# STEP 3
#========



#========
# STEP 4
#========



#========
# STEP 5
#========


#======================
# SIMULATE THE CIRCUIT
#======================



#========
# STEP 6
#========


#========
# STEP 1
#========
n = 100
# COMPLETE THIS CODE WITH YOUR SOLUTION FROM PART 2


#================
# EVE INTERCEPTS!
#================
for i in range(# COMPLETE THIS LINE
  alice_bob_qubits[i].measure(# COMPLETE THIS LINE

backend = Aer.get_backend('qasm_simulator')
job = execute(alice_bob_qubits, backend = backend, shots = 1)
result = job.result()
counts = result.get_counts()

eve_alice_bits = []
eve_bob_bits = []

for # COMPLETE THIS LINE

  # Looks at measurement results
  bits = list(counts[i].keys())[0]
  eve_alice_bits += [bits[0]]
  eve_bob_bits += [# COMPLETE THIS LINE

  # Prepares new qubits for Alice and Bob
  alice_bob_qubits[i] = QuantumCircuit(# COMPLETE THIS LINE

  # Makes sure they are in the same state she measured
  if eve_alice_bits[i] == 1: alice_bob_qubits[i].# COMPLETE THIS LINE
  if # COMPLETE THIS LINE



#========
# STEP 6
#========
alice_key = []
alice_mismatched_choices = []
alice_mismatched_choice_bits = []

bob_key = []
bob_mismatched_choices = []
bob_mismatched_choice_bits = []

eve_key = []


# COMPLETE THIS CODE


entanglement = entanglement_amount(alice_mismatched_choices, alice_mismatched_choice_bits, bob_mismatched_choices, bob_mismatched_choice_bits)
print("Entanglement of Mismatched Choices: " + str(entanglement))
print("Alice's Key: " + str(alice_key))
print("Bob's Key: " + str(bob_key))
print("Eve's Key: " + str(eve_key))
print("Key Length: " + str(len(bob_key)))
print("Number of Disagreeing Key Bits between Alice and Bob: " + str(sum([alice_key[i] != bob_key[i] for i in range(len(alice_key))])))
print("Number of Disagreeing Key Bits between Alice and Eve: " + str(sum([alice_key[i] != eve_key[i] for i in range(len(alice_key))])))
print("Number of Disagreeing Key Bits between Bob and Eve: " + str(sum([bob_key[i] != eve_key[i] for i in range(len(alice_key))])))

