#!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.h(0)
alice_bob_qubits.cx(0,1)
alice_bob_qubits.x(0)
alice_bob_qubits.z(0)

alice_bob_qubits.draw()

alice_option_1 = QuantumCircuit(1, 1)

alice_option_1.rz(0,0)

alice_option_1.h(0)

alice_option_1.measure(0, 0)

alice_option_2 = QuantumCircuit(1, 1)

alice_option_2.rz(math.pi/4,0)  # COMPLETE THIS LINE

alice_option_2.h(0)

alice_option_2.measure(0, 0)

alice_option_3 = QuantumCircuit(1, 1)

alice_option_3.rz(math.pi/2,0)   # COMPLETE THIS LINE

alice_option_3.h(0)

alice_option_3.measure(0, 0)

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

alice_circuit = alice_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.rz(math.pi/4,0)# COMPLETE THIS LINE
bob_option_1.h(0)           # COMPLETE THIS LINE

bob_option_1.measure(0, 0)

bob_option_2 = QuantumCircuit(1, 1)

bob_option_2.rz(math.pi/4,0)# COMPLETE THIS LINE

bob_option_2.h(0)           # COMPLETE THIS CODE

bob_option_2.measure(0, 0)

bob_option_3= QuantumCircuit(1, 1) # COMPLETE THIS CODE

bob_option_3.rz(math.pi*3/4,0)

bob_option_3.h(0)

bob_option_3.measure(0, 0)

bob_options = [bob_option_1,bob_option_2,bob_option_3]# COMPLETE THIS LINE
bob_choice = random.randint(0, 2)

bob_circuit = bob_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)
alice_bob_qubits.h(0)  
alice_bob_qubits.cx(0,1)
alice_bob_qubits.x(0)
alice_bob_qubits.z(0)
alice_bob_qubits.draw()



#========
# STEP 2
#========
alice_choice = random.randint(0, 2)

alice_circuit = alice_options[alice_choice]
alice_circuit.draw()

#========
# STEP 3
#========
alice_bob_qubits = alice_bob_qubits.compose(alice_circuit, qubits = [0], clbits = [0])


#========
# STEP 4
#========
bob_choice = random.randint(0, 2)
bob_circuit = bob_options[bob_choice]


#========
# STEP 5
#========
alice_bob_qubits = alice_bob_qubits.compose(bob_circuit, qubits = [1], clbits = [1])


#======================
# 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(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))

n = 100

alice_bob_qubits = []

for i in range(n):
    alice_bob_qubits += [QuantumCircuit(2,2)]

    alice_bob_qubits[i].h(0)    
    alice_bob_qubits[i].cx(0,1) 
    alice_bob_qubits[i].x(0)    
    alice_bob_qubits[i].z(0)    

alice_choices = []
alice_circuits = []

for i in range(n):                           
    alice_choices += [random.randint(0, 2)]
    alice_circuits += [alice_options[alice_choices[i]]] 

for i in range(100):                                      
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [0] , clbits = [0])


bob_choices = []
bob_circuits = []

for i in range(n):                           
    bob_choices += [random.randint(0, 2)]
    bob_circuits += [bob_options[bob_choices[i]]]

for i in range(100):
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [1] , clbits = [1])

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 i in range(n):

  # MATCHING CHOICE
  if alice_options[alice_choices[i]] == bob_options[bob_choices[i]]: 
    alice_key += [int(alice_bits[i])]
    bob_key += [1 - int(bob_bits[i])]

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

    alice_mismatched_choice_bits += [alice_bits[i]] # COMPLETE THIS LINE
    bob_mismatched_choice_bits +=  [bob_bits[i]]    # 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
#========
alice_bob_qubits = []

for i in range(n):
    alice_bob_qubits += [QuantumCircuit(2,2)]

    alice_bob_qubits[i].h(0)    
    alice_bob_qubits[i].cx(0,1) 
    alice_bob_qubits[i].x(0)    
    alice_bob_qubits[i].z(0)    


#========
# STEP 3
#========
alice_choices = []
alice_circuits = []

for i in range(n):                           
    alice_choices += [random.randint(0, 2)]
    alice_circuits += alice_options[alice_choice]       
for i in range(100):                                   
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [0] , clbits = [0])

#========
# STEP 4
#========
bob_choices = []
bob_circuits = []

for i in range(n):                           
    bob_choices += [random.randint(0, 2)]
    bob_circuits += bob_options[bob_choice]


#========
# STEP 5
#========
for i in range(100):
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [1] , clbits = [1])

#======================
# SIMULATE THE CIRCUIT
#======================
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]]


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

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

for i in range(n):

  # MATCHING CHOICE
  if alice_options[alice_choices[i]] == bob_options[bob_choices[i]]: 
    alice_key += [int(alice_bits[i])]
    bob_key += [1 - int(bob_bits[i])]

  # MISMATCHING CHOICE
  else:
    alice_mismatched_choices += [alice_choices[i]]
    bob_mismatched_choices += [bob_choices[i]] 

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


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
# COMPLETE THIS CODE WITH YOUR SOLUTION FROM PART 2


#================
# EVE INTERCEPTS!
#================
for i in range(n):
    alice_bob_qubits[i].measure(0,0)

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 i in range(100):# COMPLETE THIS LINE

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

    # Prepares new qubits for Alice and Bob
    alice_bob_qubits[i] = QuantumCircuit(2,2)

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

alice_choices = []
alice_circuits = []

for i in range(n):                           
    alice_choices += [random.randint(0, 2)]
    alice_circuits += alice_options[alice_choice]        
for i in range(100):                                     
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [0] , clbits = [0])

#========
# STEP 4
#========
bob_choices = []
bob_circuits = []

for i in range(n):                           
    bob_choices += [random.randint(0, 2)]
    bob_circuits += bob_options[bob_choice]


#========
# STEP 5
#========
for i in range(100):
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [1] , clbits = [1])

#======================
# SIMULATE THE CIRCUIT
#======================
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]]


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

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

eve_key = []
eve_mismatched_choices = []
eve_mismatched_choice_bits = []

for i in range(n):

  # MATCHING CHOICE
  if alice_options[alice_choices[i]] == bob_options[bob_choices[i]]: 
    alice_key += [int(alice_bits[i])]
    bob_key += [1 - int(bob_bits[i])]
    eve_key += [int(eve_bob_bits[i])]

 
  else:
    alice_mismatched_choices += [alice_choices[i]]
    bob_mismatched_choices += [bob_choices[i]] 
    bob_mismatched_choices += [bob_choices[i]]

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


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))])))

n = 100
# COMPLETE THIS CODE WITH YOUR SOLUTION FROM PART 2
alice_bob_qubits = []

for i in range(n):# COMPLETE THIS LINE
    alice_bob_qubits += [QuantumCircuit(2,2)]# COMPLETE THIS LINE

    alice_bob_qubits[i].h(0)    # COMPLETE THIS LINE
    alice_bob_qubits[i].cx(0,1) # COMPLETE THIS LINE
    alice_bob_qubits[i].x(0)    # COMPLETE THIS LINE
    alice_bob_qubits[i].z(0)    # COMPLETE THIS LINE


#========
# STEP 3
#========
alice_choices = []
alice_circuits = []

for i in range(n):                           # COMPLETE THIS LINE
    alice_choices += [random.randint(0, 2)]
    alice_circuits += [alice_options[alice_choices[i]]]         # COMPLETE THIS LINE
for i in range(100):                                      # COMPLETE THIS LINE
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [0] , clbits = [0])# COMPLETE THIS LINE

#========
# STEP 4
#========
bob_choices = []
bob_circuits = []

for i in range(n):                           # COMPLETE THIS LINE
    bob_choices += [random.randint(0, 2)]
    bob_circuits += [bob_options[bob_choices[i]]]# COMPLETE THIS CODE


#========
# STEP 5
#========
for i in range(100):
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [1] , clbits = [1])# COMPLETE THIS CODE


#================
# EVE INTERCEPTS!
#================
for i in range(n):# COMPLETE THIS LINE
    alice_bob_qubits[i].measure(0,0)# 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 i in range(100):# COMPLETE THIS LINE

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

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

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

#================

alice_choices = []
alice_circuits = []

for i in range(n):                           # COMPLETE THIS LINE
    alice_choices += [random.randint(0, 2)]
    alice_circuits += [alice_options[alice_choices[i]]]         # COMPLETE THIS LINE
for i in range(100):                                      # COMPLETE THIS LINE
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [0] , clbits = [0])# COMPLETE THIS LINE

#========
# STEP 4
#========
bob_choices = []
bob_circuits = []

for i in range(n):                           # COMPLETE THIS LINE
    bob_choices += [random.randint(0, 2)]
    bob_circuits += [bob_options[bob_choices[i]]]# COMPLETE THIS CODE


#========
# STEP 5
#========
for i in range(100):
    alice_bob_qubits[i] = alice_bob_qubits[i].compose(alice_circuit, qubits = [1] , clbits = [1])# COMPLETE THIS CODE

#======================
# SIMULATE THE CIRCUIT
#======================
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]]

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

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

eve_key = []
eve_mismatched_choices = []
eve_mismatched_choice_bits = []

for i in range(n):# COMPLETE THIS LINE

  # MATCHING CHOICE
  if alice_options[alice_choices[i]] == bob_options[bob_choices[i]]: # COMPLETE THIS LINE
    alice_key += [int(alice_bits[i])]
    bob_key += [1 - int(bob_bits[i])]# COMPLETE THIS LINE
    eve_key += [int(eve_bob_bits[i])]

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

    alice_mismatched_choice_bits += [alice_bits[i]] # COMPLETE THIS LINE
    bob_mismatched_choice_bits +=  [bob_bits[i]]    # COMPLETE THIS LINE# 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))])))