from PIL import Image
from math import sqrt


def newtonFunction(z):
    for i in xrange(20):
        z -= (z**3 - 1)/(3 * z**2)
    return z


def draw(size, viewport=((-10, -10), (10, 10))):
    # The complex attractors
    solutions = (complex(1), complex(-0.5, sqrt(3)/2), complex(-0.5, -sqrt(3)/2))
    img = Image.new("RGB", (size, size))

    # Calculate shifts and lengths of sides
    lens = abs(viewport[0][0] - viewport[1][0]), abs(viewport[0][1] - viewport[1][1])
    step = lens[0]/float(size), lens[1]/float(size)

    for x in xrange(size):
        for y in xrange(size):
            val = newtonFunction(complex(viewport[0][0] + x*step[0], viewport[0][1] + y*step[1]))

            # Find the closest attractor and color the pixel acordingly
            distance = map(lambda z: abs(z-val), solutions)
            color = [0, 0, 0]
            color[distance.index(min(distance))] = 255
            img.putpixel((x, y), tuple(color))

    img.save("Newton.png")

# If the newtonFunction gets 0+0i you will get devision by zero error
draw(201)
