import random
from qiskit import QuantumCircuit, QuantumRegister, execute, Aer 
from qiskit.extensions import UnitaryGate
import numpy as np

# Generate a random bit
random_bit = random.randint(0, 1)

# Print the random bit
print(random_bit)

# Charlie to Alice bit
CA = random.randint(0, 1)

# Charlie to Bob bit
CB = random.randint(0, 1)

print(CA,CB)

# Alice to Charlie bit
AC = 0

# Bob to Charlie bit
BC = 0

# Charlie checks
(CA and CB) == (AC or BC)

classical_wins = 0
number_of_rounds = 10000
for i in range(number_of_rounds):
    CA = random.randint(0, 1)
    CB = random.randint(0, 1)
    
    AC, BC = 0, 0
    if (CA and CB) == (AC or BC):
        classical_wins += 1
    
print(classical_wins/number_of_rounds)

qc = QuantumCircuit(2, 2)
qc.h(0)
qc.cx(0, 1)

# qc.measure(0,0)
# Measure both qubits and store the results in classical registers
qc.measure([0, 1], [0, 1])

# Execute the circuit on a simulator backend
backend = Aer.get_backend('qasm_simulator')
job = execute(qc, backend, shots=10)
result = job.result()

# Get the counts of each measurement outcome
counts = result.get_counts(qc)
print(counts)

CA, CB = 0,0

# Create a quantum circuit with two qubits and two classical registers
qc = QuantumCircuit(2, 2)

# Prepare an entangled state of the two qubits
qc.h(0)
qc.cx(0,1)

# Apply Alice's measurement
if CA == 1:
    qc.h(0)

Bob_0_unitary = np.array([[np.cos(np.pi/8), np.sin(np.pi/8)], [-np.sin(np.pi/8), np.cos(np.pi/8)]])     
Bob_1_unitary = np.array(
    [[np.cos(np.pi/8), -np.sin(np.pi/8)],[np.sin(np.pi/8), np.cos(np.pi/8)]]
    )    


if CB == 0:
    qc.unitary(Bob_0_unitary, [1])

if CB == 1:
    qc.unitary(Bob_1_unitary, [1])

# Measure both qubits in the computational basis
qc.measure([0, 1], [0, 1])
qc.draw('mpl')


def CHSH_strategy(CA, CB):
    # Create a quantum circuit with two qubits and two classical registers
    qc = QuantumCircuit(2, 2)
    
    # Prepare an entangled state of the two qubits
    qc.h(0)
    qc.cx(0,1)
    
    # Apply Alice's measurement
    if CA == 1:
        qc.h(0)
        
    Bob_0_unitary = np.array([[np.cos(np.pi/8), np.sin(np.pi/8)], [-np.sin(np.pi/8), np.cos(np.pi/8)]])     
    Bob_1_unitary = np.array(
        [[np.cos(np.pi/8), -np.sin(np.pi/8)],[np.sin(np.pi/8), np.cos(np.pi/8)]]
        )    
        
        
    if CB == 0:
        qc.unitary(Bob_0_unitary, [1])
        
    if CB == 1:
        qc.unitary(Bob_1_unitary, [1])
        
    # Measure both qubits in the computational basis
    qc.measure([0, 1], [0, 1])
    drawing = qc.draw('mpl')

    
    # Execute the circuit on a simulator backend
    backend = Aer.get_backend('qasm_simulator')
    job = execute(qc, backend, shots=1)
    result = job.result()

    # Get the counts of each measurement outcome
    counts = result.get_counts(qc)
    # Compute the XOR of the keys
    xor_values = [int(k[0]) ^ int(k[1]) for k in counts.keys()]
    
    return xor_values[0]

# Test the CHSH strategy
quantum_wins = 0
number_of_rounds = 100
for i in range(number_of_rounds):
    CA = random.randint(0, 1)
    CB = random.randint(0, 1)
    
    output_bit = CHSH_strategy(CA,CB)
    if (CA and CB) == output_bit:
        quantum_wins += 1
    
print(quantum_wins/number_of_rounds)