import numpy as np
from PIL import Image
from skimage import color
# --- Config ---
WIDTH, HEIGHT = 4096, 4096
TOTAL_COLORS = WIDTH * HEIGHT
# --- Step 1: Generate and perceptually sort all RGB colors ---
colors = np.array([[r, g, b] for r in range(256)
for g in range(256)
for b in range(256)], dtype=np.uint8)
# Convert to LAB color space for perceptual sorting
lab_colors = color.rgb2lab(colors.reshape((-1, 1, 3)) / 255.0).reshape(-1, 3)
sorted_indices = np.lexsort((lab_colors[:, 2], lab_colors[:, 1], lab_colors[:, 0]))
colors = colors[sorted_indices]
# --- Step 2: Generate polar coordinates and conformal compactification mask ---
yy, xx = np.mgrid[0:HEIGHT, 0:WIDTH]
cx, cy = WIDTH // 2, HEIGHT // 2
x = (xx - cx) / cx
y = (yy - cy) / cy
r = np.sqrt(x**2 + y**2)
theta = np.arctan2(y, x)
# Avoid division by zero
epsilon = 1e-9
r = np.clip(r, 0, 1 - epsilon)
# Conformal transformation
compact_r = np.tanh(2.0 * np.arctanh(r))
x_compact = compact_r * np.cos(theta)
y_compact = compact_r * np.sin(theta)
# Create Penrose mask (diamond inside a conformally compactified space)
mask = (np.abs(x_compact) + np.abs(y_compact)) <= 1.0
# Coordinate lists
coords_mask = np.stack([yy[mask], xx[mask]], axis=-1)
coords_rest = np.stack([yy[~mask], xx[~mask]], axis=-1)
all_coords = np.concatenate([coords_mask, coords_rest], axis=0)
# --- Step 3: Assign each RGB color exactly once ---
assert len(all_coords) == TOTAL_COLORS
output = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
for i, (y, x) in enumerate(all_coords):
output[y, x] = colors[i]
# --- Save output ---
Image.fromarray(output).save("penrose_rgb_compactified.png")
print("Saved as 'penrose_rgb_compactified.png'")
