import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Function to determine if a number is prime
def is_prime(n):
    if n < 2: return False
    for i in range(2, int(np.sqrt(n)) + 1):
        if n % i == 0: return False
    return True

# Generate primes and composites modulo 6
N = 1000  # Range of numbers to check
primes_mod1 = [n for n in range(5, N) if is_prime(n) and n % 6 == 1]
primes_mod5 = [n for n in range(5, N) if is_prime(n) and n % 6 == 5]
composites = [n for n in range(5, N) if not is_prime(n) and (n % 6 == 1 or n % 6 == 5)]

# Function to map numbers to the complex plane using Euler's formula
def map_to_complex_plane(n, N):
    theta = 2 * np.pi * n / N  # Angle based on the number
    return np.cos(theta), np.sin(theta)

# Prepare 3D data for plotting
x_prime1, y_prime1, z_prime1 = zip(*[(map_to_complex_plane(n, N)[0], 
                                     map_to_complex_plane(n, N)[1], n) for n in primes_mod1])
x_prime5, y_prime5, z_prime5 = zip(*[(map_to_complex_plane(n, N)[0], 
                                     map_to_complex_plane(n, N)[1], n) for n in primes_mod5])
x_comp, y_comp, z_comp = zip(*[(map_to_complex_plane(n, N)[0], 
                               map_to_complex_plane(n, N)[1], n) for n in composites])

# Create the 3D plot
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

# Plot primes (1 mod 6)
ax.scatter(x_prime1, y_prime1, z_prime1, color='blue', label='Primes ≡ 1 mod 6', s=10)

# Plot primes (5 mod 6)
ax.scatter(x_prime5, y_prime5, z_prime5, color='red', label='Primes ≡ 5 mod 6', s=10)

# Plot composites
ax.scatter(x_comp, y_comp, z_comp, color='orange', label='Composites (1,5 mod 6)', alpha=0.6, s=10)

# Critical Line (z-axis)
ax.plot([0, 0], [0, 0], [0, N], color='green', linestyle='--', linewidth=2, label='Critical Line (z-axis)')

# Draw lines connecting primes to the critical line center (0, 0, z)
for i in range(len(primes_mod1)):
    ax.plot([0, x_prime1[i]], [0, y_prime1[i]], [z_prime1[i], z_prime1[i]], color='blue', alpha=0.5)
for i in range(len(primes_mod5)):
    ax.plot([0, x_prime5[i]], [0, y_prime5[i]], [z_prime5[i], z_prime5[i]], color='red', alpha=0.5)

# Draw vertical lines upward for composites
for i in range(len(composites)):
    ax.plot([x_comp[i], x_comp[i]], [y_comp[i], y_comp[i]], [0, z_comp[i]], color='orange', alpha=0.5, linestyle=':')

# Set labels and legend
ax.set_title('3D Distribution of Primes and Composites Modulo 6 with Connections to Critical Line')
ax.set_xlabel('Re(e^iθ) [Cos(θ)]')
ax.set_ylabel('Im(e^iθ) [Sin(θ)]')
ax.set_zlabel('Number (z)')
ax.legend()
plt.show()