import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm


def circle_sdf(x, y, radius, center_x, center_y):
    """
    Signed Distance Function for a 2D circle
    """
    return np.sqrt((x - center_x)**2 + (y - center_y)**2) - radius

def union_sdf(x, y, circles):
    """
    Union of multiple circle SDFs using min operation
    """
    sdfs = [circle_sdf(x, y, radius, cx, cy) for (cx, cy, radius) in circles]
    return np.min(sdfs, axis=0)

# Create grid
x = np.linspace(-3, 3, 300)
y = np.linspace(-3, 3, 300)
X, Y = np.meshgrid(x, y)

# Define circles: (center_x, center_y, radius)
circles = [
    (0, 0, 1),     # Center circle
    (1.5, 1, 0.8), # Top right circle
    (-1.5, -1, 0.6), # Bottom left circle
    (0.5, -1.5, 0.7) # Bottom right circle
]

# Calculate SDF
Z = union_sdf(X, Y, circles)

# Visualization
plt.figure(figsize=(6, 5))

# SDF Plot
# Create a normalization that centers white at zero
norm = TwoSlopeNorm(vmin=Z.min(), vcenter=0, vmax=Z.max())
contour = plt.contourf(X, Y, Z, levels=20, cmap='coolwarm', norm=norm)
plt.contour(X, Y, Z, levels=20, colors='black', linewidths=0.2)
plt.colorbar(contour)#, label='Signed Distance')
#plt.title('Signed Distance Function')
plt.xlabel('$x$')
plt.ylabel('$y$')

# Add circle centers
for cx, cy, r in circles:
    plt.plot(cx, cy, 'ko')
    circle = plt.Circle((cx, cy), r, fill=False, color='black')
    plt.gca().add_artist(circle)

plt.tight_layout()

if True:
    plt.savefig('signed-distance-function.png', dpi=300)

plt.show()