import math
from PIL import Image

width, height = 800, 800
max_iter = 40
tolerance = 0.001  # Wie nah müssen wir an der Nullstelle sein?

# Die drei Nullstellen von z^3 - 1 = 0
roots = [
    complex(1, 0),
    complex(-0.5, math.sqrt(3) / 2),
    complex(-0.5, -math.sqrt(3) / 2)
]

# Farben für die drei Nullstellen (Rot, Grün, Blau)
root_colors = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255)
]

img = Image.new('RGB', (width, height), (0, 0, 0))
pixels = img.load()

for py in range(height):
    for px in range(width):
        # Koordinaten umrechnen (von -2 bis 2)
        zx = -2.0 + (px / width) * 4.0
        zy = -2.0 + (py / height) * 4.0
        z = complex(zx, zy)

        found = False
        for n in range(max_iter):
            if abs(z) < 0.000001: break  # Vermeidung von Division durch Null

            # Newton-Schritt für f(z) = z^3 - 1
            # Formel: z = z - f(z) / f'(z)  =>  z = z - (z^3 - 1) / (3 * z^2)
            z_next = z - (z ** 3 - 1) / (3 * z ** 2)

            # Prüfen, ob wir nah genug an einer der Wurzeln sind
            for i, root in enumerate(roots):
                diff = z_next - root
                if abs(diff) < tolerance:
                    # Farbe basierend auf der Wurzel, Helligkeit basierend auf Iterationen
                    brightness = int(255 * (1 - n / max_iter))
                    r, g, b = root_colors[i]
                    pixels[px, py] = (
                        int(r * (brightness / 255)),
                        int(g * (brightness / 255)),
                        int(b * (brightness / 255))
                    )
                    found = True
                    break
            if found: break
            z = z_next

img.save("newton_fractal.png")
img.show()