#%matplotlib inline

from qiskit import *
from qiskit.tools.monitor import *
from qiskit.providers.ibmq import least_busy
import numpy as np

#Function to create the circuits for the CHSH game

def CHSH_circuit(x,y,a0=0,a1=np.pi/2,b0=np.pi/4,b1=-np.pi/4):
    
    #x: bit received by Alice
    #y: bit received by Bob
    #a0: measure angle used by Alice when she receives 0
    #a1: measure angle used by Alice when she receives 1
    #b0: measure angle used by Bob when he receives 0
    #b1: measure angle used by Bob when he receives 1
    
    circ = QuantumCircuit(2,2) 
    
    # First, we create a Bell pair
    
    circ.h(0)
    circ.cx(0,1)

    # Now, we apply rotations for Alice and Bob depending on the bits they have received
    
    if(x==0):
        circ.ry(a0,0)
    else:
        circ.ry(a1,0)

    if(y==0):
        circ.ry(b0,1)
    else:
        circ.ry(b1,1)

    # We measure
        
    circ.measure(range(2),range(2)) # Medimos

    return circ
    

def winning_probability(backend, shots = 8192, a0=0,a1=np.pi/2,b0=np.pi/4,b1=-np.pi/4):
    total = 0
    circuits = [CHSH_circuit(0,0,a0,a1,b0,b1), CHSH_circuit(0,1,a0,a1,b0,b1), CHSH_circuit(1,0,a0,a1,b0,b1), CHSH_circuit(1,1,a0,a1,b0,b1)] # We 'pack' four different circuits for execution
    job = execute(circuits, backend=backend, shots = shots)
    # For the first three circuits, the winning condition is that Alice's and Bob's outputs are equal
    for qc in circuits[0:3]:
        counts = job.result().get_counts(qc)
        if('00' in counts):
            total += counts['00']
        if('11' in counts):
            total += counts['11']
    # For the fourth circuit, Alice's and Bob's outputs must be different for them to win
    counts = job.result().get_counts(circuits[3])
    if('01' in counts):
        total += counts['01']
    if('10' in counts):
        total += counts['10']
        
    return total/(4*shots)    

# Backend definition 

backend = Aer.get_backend('qasm_simulator')

# Execution

print(winning_probability(backend))

# Load the account

provider = IBMQ.load_account()

# We execute on the least busy device (among the actual quantum computers)

backend = least_busy(provider.backends(operational = True, simulator=False, status_msg='active',
                                       filters=lambda x: x.configuration().n_qubits > 1)) 
print("We are executing on...",backend)

print(winning_probability(backend))