#!/usr/bin/env python

import numpy as np
from PIL import Image

# Some parameters
pixels      = (1000, 1000)
x_lims      = (-2.0, 2.0)
y_lims      = (-2.0, 2.0)
max_iter    = 16
max_value   = 2

def f(z, c):
    # Function which does the iterations
    # Takes z_{n} as input, returns z_{n+1}
    return z**2 + c

def draw_picture(picture, filename):
    # Save the image to file
    im = Image.fromarray(picture, 'L')
    im.save(filename)

def main():
    # Setting up the area to calculate for
    # Create two vectors, the number of points is given by the variable pixels,
    # the points are equally spaced in the ranges given by x_lims and y_lims
    x_range = np.linspace(x_lims[0], x_lims[1], pixels[0])
    y_range = np.linspace(y_lims[0], y_lims[1], pixels[1])

    # This line creates a 2D grid of complex numbers of the size given by
    # the variable pixels, and the ranges given by x_lims and y_lims
    C = np.meshgrid(y_range, x_range)[0] + 1j*np.meshgrid(y_range, x_range)[1]

    # Array to hold temporary values of z. The array is the same size as C, 
    # and initially filled with zeros
    Z = np.zeros(pixels, dtype = np.complex128)
    
    # Do a number of iterations
    for i in xrange(max_iter):
        # Calculate a new value of z based on the previous value, using
        # the function f(z, c)
        Z = f(Z, C)

    # Free up some memory
    C = 0

    # Some of the numbers will now be too large to store correctly in memory,
    # and these will be stored as NaN. Since we are only interested in whether
    # they are greater or smaller than 2, we set these to a large number
    Z[np.isnan(Z)] = 1e10

    # Create an array to hold the picture data
    # The type uint8 is suitable for saving as an 8 bit black and white picture
    picture = np.zeros(pixels, dtype = np.uint8)

    # Calculate the absolute value, and set those pixels which are outside the
    # set to white (255). The rest remains black (0).
    picture[abs(Z) > max_value] = 255

    # Draw and save picture
    draw_picture(picture, "mandelbrot.png")

if __name__ == '__main__':
    main()