from pylab import *
from scipy import weave

# The C++ code is written in a multiline string.
# It will be compiler, when run for the first time. Next time,
# it will just use the "old" executable.
code="""
     using namespace std;
     
     return_val=0;
     for (int i=0; i<max_iterations; i++){
         if (norm(z)>4.) { return_val= 1./i; break; }
         z = z*z + z0;     // if |z|>2 the point is not part of mandelbrot set
     }
     """

Nx = 400
Ny = 400
data = zeros((Nx,Ny)) # allocate all arrays in Python!
max_iterations=1000
ext=[-1.8,-1.72,-0.05,0.05] # The range of the mandelbrot plot [x0,x1,y0,y1]


# The double loop is written in Python, which takes some time.
# Not the most optimal code.
for i in range(Nx):
    for j in range(Ny):
        z0 = ext[0] + (ext[1]-ext[0])*i/(Nx-1.) + 1j*(ext[2] + (ext[3]-ext[2])*j/(Ny-1.) )
        z = 0j
        # This line compiles and executed the code.
        # First argument - the code, second - all necessary variables, the rest - options.
        data[i,j] = weave.inline(code, ['max_iterations', 'z0', 'z'],
                                 type_converters=weave.converters.blitz, compiler = 'gcc')

data = data.transpose() # for plotting, need to transpose

# Using python's pylab, we display pixels to the screen!
imshow(data, interpolation='bilinear', cmap=cm.hot, origin='lower', extent=ext, aspect=1.)
colorbar()
show()