#!/usr/bin/env python
from matplotlib import *
from pylab import *
import mandl

def Mandel(x, y, niter):
    "Generates Mandelbrot set for given numbers x and y"
    z0=complex(x,y)
    z=0
    for i in range(niter):
        z = z**2 + z0
        if (abs(z)>2.): return 1/(i+1.)
    return 0


def Cmp_Mandel(x, y, niter=30):
    "Generates Mandelbrot set for given array or real numbers x and y"
    Z=zeros((len(y),len(x)),dtype=float) # creates 2D-array of zeros

    for iy in range(len(y)):
        for ix in range(len(x)):
            Z[iy,ix] = Mandel(x[ix],y[iy], niter)
    return Z

xmin = -2; xmax = 1
ymin = -1; ymax = 1
dh = 0.001

x = arange(xmin,xmax,dh) # creates equidistant one dimensional array
y = arange(ymin,ymax,dh)

# Call to python function
#Z = Cmp_Mandel(x,y)

# Call to equivalent fortran function
Z = transpose(mandl.cmp_mandel(x,y,100))
#Z = mandl.cmp_mandel(x,y)


# Below is plotting
# creates nice color-maps
maps=[m for m in cm.datad.keys() if not m.endswith("_r")]
# plots the data
icolor = 12 # can be changed and colors will change
im = imshow(Z, cmap = cm.get_cmap(maps[icolor],40), extent=(xmin,xmax,ymin,ymax))

#savefig('mandel_plot') # saves the plot in mandel_plot.png
show() # shows the plot on the screen