#!/usr/bin/env python

import sys
from Numeric import arange
from PIL import Image

def r_sq(z):
    """
        Returns a square of a modulus of a complex number.
    """
    return z.real**2 + z.imag**2

def mandelbrot(c, z=0+0j, i=0):
    """
        Checks if the number escapes to infinity
    """
    if r_sq(z) > 4:
        return i
    elif i == max_i:
        return -1
    else:
        return mandelbrot(c, (z**2 + c), i+1)

def get_color(iterations):
    """
        Returns a color for a pixel, depending on when
        the number escaped (if at all) to infinity.
    """

    #looks like the scheme is BGR, not RGB
    index = {
        (0, 1):      0xff0000,
        (1, 2):      0xff1000,
        (2, 3):      0xff2000,
        (3, 4):      0xff3000,
        (4, 5):      0xff4000,
        (5, 6):      0xff5000,
        (6, 7):      0xff6000,
        (7, 8):      0xff7000,
        (8, 10):     0xff9000,
        (10, 15):    0xffa600,
        (15, 21):    0xffb500,
        (21, 50):    0xffdf00,
        (50, max_i+1): 0xffffff,
    }
    for point in index:
        if iterations in range(point[0], point[1]):
            return index[point]


if __name__ == "__main__":

    if len(sys.argv) != 2:
        print "Usage:\n\t%s <width>\n" % sys.argv[0]
        sys.exit()

    #how many iterations will be done before..
    # ..assuming that a number did not escape
    max_i = 200

    #relative proportions
    r_width = 3.2
    r_height = 2.4

    #ratio used to determine height. (avoids distortion)
    ratio = r_width / r_height

    width = int(sys.argv[1])
    height = int(width / ratio)

    #used to determine complex number for a specific pixel
    x_step = r_width / width
    y_step = r_height / height

    #real axis
    x_axis = arange(-2.2, 1.0, x_step).tolist()
    #imaginary axis
    y_axis = arange(-1.2, 1.2, y_step).tolist()

    out = []

    for y in y_axis:
        for x in x_axis:
            result = mandelbrot(complex(x, y))
            if result < 0:
                #points that did not escape
                out.append(0x000000)
            else:
                #points that did escape
                out.append(get_color(result))

    img = Image.new('RGB', (width, height))
    img.putdata(out)
    img.save('m.png')