import matplotlib.pyplot as plt
import numpy as np

def khintchine_experiment(N, num_trials=1000):
  complex_numbers = np.random.randn(N) + 1j * np.random.randn(N)

  sums = []
  for _ in range(num_trials):
    # Generate random signs (+1 or -1)
    signs = np.random.choice([-1, 1], size=N)

    # Multiply complex numbers by random signs and sum them
    current_sum = np.sum(complex_numbers * signs)
    sums.append(current_sum)

  # Calculate the expected modulus
  expected_modulus = np.sqrt(np.sum(np.abs(complex_numbers)**2))

  return complex_numbers, sums, expected_modulus

# Parameters for the experiment
N = 10  # Number of complex numbers
num_trials = 3000

# Run the experiment
complex_numbers, sums, expected_modulus = khintchine_experiment(N, num_trials)

# Plotting
plt.figure(figsize=(8, 8))

# Plot the original complex numbers
plt.scatter(
    [x.real for x in complex_numbers],
    [x.imag for x in complex_numbers],
    color='blue',
    label='$x_i$',
    marker='x'
)

# Plot the resulting sums
plt.scatter(
    [s.real for s in sums],
    [s.imag for s in sums],
    color='black',
    alpha=0.1,
    s=5,
    label='$\sum \epsilon_i x_i{{))}} 
)

# Plot circles representing the expected modulus range
circle_avg = plt.Circle((0, 0), np.mean(np.abs(sums)), color='red', fill=False, linestyle='--', label=f'$\\mathbb{{E{{))}}\\mid\\sum_i \epsilon_i x_i\\mid{{))}})
circle_expected = plt.Circle((0, 0), expected_modulus, color='purple', fill=False, linestyle='-.', label=f'$\sqrt{{\sum_i \\mid x_i \\mid^2{{))}})
plt.gca().add_patch(circle_avg)
plt.gca().add_patch(circle_expected)

plt.title(f'Khintchine Inequality (N={N}, Trials={num_trials})')
plt.legend()
plt.xticks([]), plt.yticks([])
plt.axis('equal')  # Ensure circles are displayed as circles
plt.show()