import numpy as np
import pandas as pd
import re
import random
import math

# Importing QISKit
from qiskit import QuantumProgram
import Qconfig

# Quantum program setup
Q_program = QuantumProgram()
Q_program.set_api(Qconfig.APItoken, Qconfig.config['url']) # set the APIToken and API url

## Input parameters

measurementChoicesLength = 1000 # length of the strings that cointain measurement choices
evePresence = False # Presence of an eavesdropper in the channel
evePresencePercentage = 1 # Percentage of Eve's interference
backend = 'local_qasm_simulator'

## Creating registers
qr = Q_program.create_quantum_register("qr", 2)
cr = Q_program.create_classical_register("cr", 4)

## Creating a shared entangled state

entangledState = Q_program.create_circuit('entangledState', [qr], [cr])

# bell_11 (singlet) stete for the ibmqx2 backend
#entangledState.x(qr[0])
#entangledState.x(qr[1])
#entangledState.h(qr[0])
#entangledState.cx(qr[0],qr[1])

# bell_11 (singlet) stete for the ibmqx4 backend
entangledState.x(qr[0])
entangledState.x(qr[1])
entangledState.h(qr[1])
entangledState.cx(qr[1],qr[0])

### Creating measurement circuits

## Alice's measurement circuits

# measurement of the spin projection onto the A1 = (1; 0; 0) direction
# projection onto the eigenvectors of the matrix X
measureA1 = Q_program.create_circuit('measureA1', [qr], [cr])
measureA1.h(qr[0])
measureA1.measure(qr[0],cr[0])

# measurement of the spin projection onto A2 = (1/sqt(2); 0; 1/(sqrt(2)) direction
# projection onto the eigenvectors of the matrix W = (Z+X)/sqrt(2)
measureA2 = Q_program.create_circuit('measureA2', [qr], [cr])
measureA2.s(qr[0])
measureA2.h(qr[0])
measureA2.t(qr[0])
measureA2.h(qr[0])
measureA2.measure(qr[0],cr[0])

# measurement of the spin projection onto the A3 = (0; 0; 1) vector
# standard Z-measurement
measureA3 = Q_program.create_circuit('measureA3', [qr], [cr])
measureA3.measure(qr[0],cr[0])

## Bob's measurement circuits

# measurement of the spin projection onto the B1 = (1/sqt(2); 0; 1/(sqrt(2)) direction
# projection onto the eigenvectors of the matrix W = (Z+X)/sqrt(2)
measureB1 = Q_program.create_circuit('measureB1', [qr], [cr])
measureB1.s(qr[1])
measureB1.h(qr[1])
measureB1.t(qr[1])
measureB1.h(qr[1])
measureB1.measure(qr[1],cr[1])

# measurement of the spin projection onto the B2 = (0; 0; 1) direction 
# standard Z measurement
measureB2 = Q_program.create_circuit('measureB2', [qr], [cr])
measureB2.measure(qr[1],cr[1])

# measurement of the spin projection onto the B3 = (-1/sqt(2); 0; 1/(sqrt(2)) direction
# projection onto the eigenvectors of the matrix V = (Z-X)/sqrt(2)
measureB3 = Q_program.create_circuit('measureB3', [qr], [cr])
measureB3.s(qr[1])
measureB3.h(qr[1])
measureB3.tdg(qr[1])
measureB3.h(qr[1])
measureB3.measure(qr[1],cr[1])

## Eve's measuremets

# identity gate for Alice's qubit 
ident0 = Q_program.create_circuit('ident0', [qr], [cr])
ident0.iden(qr[0])

# identity gate for Bob's qubit 
ident1 = Q_program.create_circuit('ident1', [qr], [cr])
ident1.iden(qr[1])

# measurement of the spin projection of Alice's qubit onto the (1; 0; 0) direction (observable X)
measureEA1 = Q_program.create_circuit('measureEA1', [qr], [cr])
measureEA1.h(qr[0])
measureEA1.measure(qr[0],cr[2])

# measurement of the spin projection of Alice's qubit onto the (1/sqt(2); 0; 1/(sqrt(2)) direction (observable W)
measureEA2 = Q_program.create_circuit('measureEA2', [qr], [cr])
measureEA2.s(qr[0])
measureEA2.h(qr[0])
measureEA2.t(qr[0])
measureEA2.h(qr[0])
measureEA2.measure(qr[0],cr[2])

# measurement of the spin projection of Alice's qubit onto the (0; 0; 1) direction (observable Z)
measureEA3 = Q_program.create_circuit('measureEA3', [qr], [cr])
measureEA3.measure(qr[0],cr[2])

# measurement of the spin projection of Bob's qubit onto the (1/sqt(2); 0; 1/(sqrt(2)) direction (observable W)
measureEB1 = Q_program.create_circuit('measureEB1', [qr], [cr])
measureEB1.s(qr[1])
measureEB1.h(qr[1])
measureEB1.t(qr[1])
measureEB1.h(qr[1])
measureEB1.measure(qr[1],cr[3])

# measurement of the spin projection of Bob's qubit onto the (0; 0; 1) direction (observable Z)
measureEB2 = Q_program.create_circuit('measureEB2', [qr], [cr])
measureEB2.measure(qr[1],cr[3])

# measurement of the spin projection of Bob's qubit onto the (-1/sqt(2); 0; 1/(sqrt(2)) direction (observable V)
measureEB3 = Q_program.create_circuit('measureEB3', [qr], [cr])
measureEB3.s(qr[1])
measureEB3.h(qr[1])
measureEB3.tdg(qr[1])
measureEB3.h(qr[1])
measureEB3.measure(qr[1],cr[3])

## Lists of Alice's, Bob's and Eve's measurement circuits

aliceMeasurements = [measureA1, measureA2, measureA3]
bobMeasurements = [measureB1, measureB2, measureB3]
eveMeasurements = [ident0, ident1, measureEA1, measureEA2, measureEA3, measureEB1, measureEB2, measureEB3]

## Alice's and Bob's measurement choice strings

aliceMeasurementChoices = []
bobMeasurementChoices = []

# random strings generation
for i in range(measurementChoicesLength):
    aliceMeasurementChoices.append(random.randint(1, 3))
    bobMeasurementChoices.append(random.randint(1, 3))

## Eve's optimal measurement choices array

# it is efficient for Eve to measure the spin projections of Alice's and Bob's qubits onto the same direction (A2B1 and A3B2) 

# list of Eve's measurement choices
# the first and the second element of each row represents the measurement of Alice's and Bob's qubit, respectively 
# default values of the list are 1 and 2 (applying identity gate)
eveMeasurementChoices = [[1,2] for j in range(measurementChoicesLength)]

eveMersurementPerformed = [0] * measurementChoicesLength

if evePresence == True:
    for j in range(measurementChoicesLength):
        if random.uniform(0, 1) <= evePresencePercentage:
            
            eveMersurementPerformed[j] = 1
            
            if random.uniform(0, 1) <= 0.5: # in 50% of cases perform the A2_B1 measurement (observable WW)
                eveMeasurementChoices[j] = [4, 6]
                
            else: # in 50% of cases perform the A3_B2 measurement (observable ZZ)
                eveMeasurementChoices[j] = [5, 7]

circuits = [] # the list in which the created circuits will be stored

for k in range(measurementChoicesLength):
    # create the name of the k-th circuit depending on Alice's, Bob's and Eve's choices of measurement
    circuitName = str(k) + ':A' + str(aliceMeasurementChoices[k]) + '_B' + str(bobMeasurementChoices[k]) + '_E' + str(eveMeasurementChoices[k][0]) + str(eveMeasurementChoices[k][1])
    
    # create the joint measurement circuit
    # add Alice's and Bob's measurement circuits to the singlet state curcuit
    Q_program.add_circuit(circuitName,
                          entangledState + # singlet state circuit
                          eveMeasurements[eveMeasurementChoices[k][0]-1] + # Eve's measurement circuit for Alice's qubit
                          eveMeasurements[eveMeasurementChoices[k][1]-1] + # Eve's measurement circuit for Bob's qubit
                          aliceMeasurements[aliceMeasurementChoices[k]-1] + # measurement circuit of Alice
                          bobMeasurements[bobMeasurementChoices[k]-1] # measurement circuit of Bob
                         )
    
    # add the created circuit to the circuits list
    circuits.append(circuitName)

## Execute circuits

result = Q_program.execute(circuits, backend=backend, shots=1, timeout=180)
print(result)

## Recording the measurement results

aliceResults = [] # Alice's results
bobResults = [] # Bob's results
eveResults = [[0, 0] for j in range(measurementChoicesLength)] # list of zeros by default; if the element is 0, then Eve did not perform a measurement

# Alice's and Bob's search patterns
abPatterns = [
    re.compile('..00$'), # search for the '..00' output (Alice obtained -1 and Bob obtained -1)
    re.compile('..01$'), # search for the '..01' output
    re.compile('..10$'), # search for the '..10' output (Alice obtained -1 and Bob obtained 1)
    re.compile('..11$')  # search for the '..11' output
]

# Eve's search patterns
ePatterns = [
    re.compile('00..$'), # search for the '00..' result
    re.compile('01..$'),
    re.compile('10..$'),
    re.compile('11..$')
]

# Recording results
for k in range(measurementChoicesLength):
    
    res = list(result.get_counts(circuits[k]).keys())[0] # extract a key from the dict and transform it to str; execution result of the k-th circuit
    
    # Alice and Bob
    if abPatterns[0].search(res): # check if the key is '..00' (if the measurement results are -1,-1)
        aliceResults.append(-1) # Alice got the result -1 
        bobResults.append(-1) # Bob got the result -1
    if abPatterns[1].search(res):
        aliceResults.append(1)
        bobResults.append(-1)
    if abPatterns[2].search(res):
        aliceResults.append(-1) 
        bobResults.append(1)
    if abPatterns[3].search(res): # check if the key is '..11' (if the measurement results are 1,1)
        aliceResults.append(1)
        bobResults.append(1)
        
    # Eve
    if eveMersurementPerformed[k] == True: # if Eve has performed a measurement in the k-th circuit
        if ePatterns[0].search(res): # check if the key is '00..'
            eveResults[k][0] = -1 # result of the measurement of Alice's qubit
            eveResults[k][1] = -1 # result of the measurement of Bob's qubit
        if ePatterns[1].search(res):
            eveResults[k][0] = 1
            eveResults[k][1] = -1
        if ePatterns[2].search(res):
            eveResults[k][0] = -1
            eveResults[k][1] = 1
        if ePatterns[3].search(res):
            eveResults[k][0] = 1
            eveResults[k][1] = 1

## Measurement choises check stage

# Alice and Bob compare their strings with basis choices
# If in the k-th measurement A and B used the same basis, then they record the result of the k-th measurement as the bit of the secret key 

aliceKey = [] # Alice's key string
bobKey = [] # Bob's key string
eveKeys = [] # Eve's keys; the 1-st column is the Alice's key, the 2-nd is Bob's 

# the key consists of the results obtained after projecting the spin onto the same directions (A2_B1 and A3_B2 measurements)
for k in range(measurementChoicesLength):
    
    # if Alice and Bob measured the WW or ZZ observable
    if (aliceMeasurementChoices[k] == 2 and bobMeasurementChoices[k] == 1) or (aliceMeasurementChoices[k] == 3 and bobMeasurementChoices[k] == 2):
        aliceKey.append(aliceResults[k]) # record Alice's k-th result as a key bit
        bobKey.append(-bobResults[k]) # write Bob's k-th result as a key bit; must be inversed for the singlet state
        eveKeys.append([eveResults[k][0], -eveResults[k][1]])
        
keyLength = len(aliceKey) # length of the secret keys

## Comparing the bits of the keys

abKeyMismatches = 0 # number of mismatching bits in Alice's and Bob's secret keys
eaKeyMismatches = 0 # number of mismatching bits in Alice's and Eve's secret keys
ebKeyMismatches = 0
eeKeyMismatches = 0

for k in range(keyLength):
    if aliceKey[k] != bobKey[k]:
        abKeyMismatches += 1
    if eveKeys[k][0] != aliceKey[k]:
        eaKeyMismatches += 1
    if eveKeys[k][1] != bobKey[k]:
        ebKeyMismatches += 1
    if eveKeys[k][0] != eveKeys[k][1]:
        eeKeyMismatches += 1

# Eve's knowledge of the keys
eaKnowledge = (keyLength - eaKeyMismatches)/keyLength # Eve's knowledge of Alice's key
ebKnowledge = (keyLength - ebKeyMismatches)/keyLength # Eve's knowledge of Bob's key
averageEveKnowledge = (eaKnowledge + ebKnowledge)/2 # average Eve's knowledge

# Quantum bit error rate
qber = abKeyMismatches/keyLength

## CHSH inequality test

# lists with the counts of measurement results
# each element represents the number of (-1,-1), (-1,1), (1,-1) and (1,1) results respectively
countA1B1 = [0, 0, 0, 0]
countA1B3 = [0, 0, 0, 0]
countA3B1 = [0, 0, 0, 0]
countA3B3 = [0, 0, 0, 0]

for k in range(measurementChoicesLength):

    res = list(result.get_counts(circuits[k]).keys())[0]
    
    # observable XW
    if (aliceMeasurementChoices[k] == 1 and bobMeasurementChoices[k] == 1):
        for j in range(4):
            if abPatterns[j].search(res):
                countA1B1[j] += 1
    
    # observable XV
    if (aliceMeasurementChoices[k] == 1 and bobMeasurementChoices[k] == 3):
        for j in range(4):
            if abPatterns[j].search(res):
                countA1B3[j] += 1
    
    # observable ZW
    if (aliceMeasurementChoices[k] == 3 and bobMeasurementChoices[k] == 1):
        for j in range(4):
            if abPatterns[j].search(res):
                countA3B1[j] += 1
    
    # observable ZV
    if (aliceMeasurementChoices[k] == 3 and bobMeasurementChoices[k] == 3):
        for j in range(4):
            if abPatterns[j].search(res):
                countA3B3[j] += 1

# number of results obtained from measurements in a particular basis
total11 = sum(countA1B1)
total13 = sum(countA1B3)
total31 = sum(countA3B1)
total33 = sum(countA3B3)                  
                
# expected values of the spin projections onto the A1_B1, A1_B3, A3_B1 and A3_B3  directions
expect11 = (countA1B1[0] - countA1B1[1] - countA1B1[2] + countA1B1[3])/total11 # observable XW
expect13 = (countA1B3[0] - countA1B3[1] - countA1B3[2] + countA1B3[3])/total31 # observable XV
expect31 = (countA3B1[0] - countA3B1[1] - countA3B1[2] + countA3B1[3])/total13 # observable ZW
expect33 = (countA3B3[0] - countA3B3[1] - countA3B3[2] + countA3B3[3])/total33 # observable ZV

# CHSH correlation value
# must be equal to -2*sqrt(2) in the case of the absence of noise and eavesdropping
corr = expect11 - expect13 + expect31 + expect33
diff = abs(abs(corr) - 2*math.sqrt(2))

## Print results

# CHSH inequality test
print('CHSH correlation value: ' + str(round(corr, 3)))
print('Difference from -2*sqrt(2): ' + str(round(diff, 3)) + '\n')

# Secret keys
print('Length of the secret key: ' + str(keyLength))
print('Number of mismatching bits: ' + str(abKeyMismatches) + '\n')

print('QBER: ' + str(round(qber * 100, 2)) + ' %')
print('Eve\'s average knowledge: ' + str(round(averageEveKnowledge * 100, 2)) + ' %\n')

## Check measurement errors

# Theoretically, the measurement in the A2_B1 and A3_B2 basises gives only the (-1,1) and (1,-1) results
# This block counts how many (-1,-1) and (1,1) results obtained after the measurements of the WW and ZZ observables
# actual for the execution on the real quantum devices

mismatchesList = []

for k in range(measurementChoicesLength):

    res = list(result.get_counts(circuits[k]).keys())[0]
    
    result01 = False # by default
    result10 = False
    
    if abPatterns[1].search(res):
        result01 = True
    if abPatterns[2].search(res):
        result10 = True
    
    condition1 = ('A2_B1' in circuits[k] or 'A3_B2' in circuits[k]) and not result01 # if result is not (-1,1)
    condition2 = ('A2_B1' in circuits[k] or 'A3_B2' in circuits[k]) and not result10 # if result is not (1,-1)
    
    if condition1 and condition2 == True:
        mismatchesList.append([str(circuits[k]), res])
        
pd.DataFrame(mismatchesList,columns = ['Measurement', 'Result'])