import math
import random

def chsh(correlations):
    if len(correlations) != 4:
        raise ValueError("CHSH needs four setting-pair correlations")
    return abs(correlations[0] + correlations[1] + correlations[2] - correlations[3])

def singlet_correlation(alice_angle, bob_angle):
    return -math.cos(alice_angle - bob_angle)

def sample_pair(correlation, shots, seed):
    if shots < 1 or not -1 <= correlation <= 1:
        raise ValueError("invalid sampling request")
    rng = random.Random(seed)
    products = []
    alice_total = bob_total = 0
    for _ in range(shots):
        alice = 1 if rng.randrange(2) else -1
        same = rng.random() < (1 + correlation) / 2
        bob = alice if same else -alice
        products.append(alice * bob)
        alice_total += alice
        bob_total += bob
    estimate = sum(products) / shots
    variance = max(0.0, 1.0 - estimate ** 2)
    return estimate, math.sqrt(variance / shots), alice_total / shots, bob_total / shots

local_values = []
for encoded in range(16):
    outputs = tuple(1 if encoded & (1 << bit) else -1 for bit in range(4))
    alice, alice_prime, bob, bob_prime = outputs
    value = abs(alice * bob + alice * bob_prime + alice_prime * bob - alice_prime * bob_prime)
    local_values.append(value)
assert max(local_values) == 2 and min(local_values) == 2

alice_angles = (0.0, math.pi / 2)
bob_angles = (math.pi / 4, -math.pi / 4)
setting_pairs = ((alice_angles[0], bob_angles[0]),
                 (alice_angles[0], bob_angles[1]),
                 (alice_angles[1], bob_angles[0]),
                 (alice_angles[1], bob_angles[1]))
ideal_correlations = tuple(singlet_correlation(*settings) for settings in setting_pairs)
ideal_s = chsh(ideal_correlations)
assert abs(ideal_s - 2 * math.sqrt(2)) < 1e-12

shots = 25000
sampled = [sample_pair(correlation, shots, 2300 + index)
           for index, correlation in enumerate(ideal_correlations)]
sampled_correlations = tuple(row[0] for row in sampled)
sampled_s = chsh(sampled_correlations)
sampled_sigma = math.sqrt(sum(row[1] ** 2 for row in sampled))
max_marginal = max(abs(value) for row in sampled for value in row[2:])
assert abs(sampled_s - ideal_s) < 6 * sampled_sigma
assert sampled_s - 2 > 8 * sampled_sigma
assert max_marginal < 0.03

local_sampled = [sample_pair(1.0, 4000, 2390 + index) for index in range(4)]
local_sampled_s = chsh(tuple(row[0] for row in local_sampled))
assert local_sampled_s == 2.0

wrong_sign_s = abs(sum(ideal_correlations))
assert wrong_sign_s <= 2 and abs(wrong_sign_s - ideal_s) > 1.0
biased_marginal = 1.0
assert biased_marginal > 5 / math.sqrt(shots)
zero_shots_rejected = False
try:
    sample_pair(0.0, 0, 1)
except ValueError:
    zero_shots_rejected = True
assert zero_shots_rejected
print(f"PASS: 23 CHSH notebook bounds all {len(local_values)} deterministic local strategies at 2, gives ideal S={ideal_s:.6f}, sampled S={sampled_s:.4f}±{sampled_sigma:.4f}, max marginal={max_marginal:.4f}")