#!/usr/bin/env python
from scipy import *
from pylab import *
from matplotlib.patches import Rectangle
import weave
## Kristjan Haule, Jan 2014

# The C++ code
code="""
    #line 10 "mand3.py"
     using namespace std; // for using cout when debugging
     for (int i=0; i<Nx; i++){
         for (int j=0; j<Ny; j++){
             complex<double> z0( ext(0)+(ext(1)-ext(0))*i/(Nx-1.), ext(2)+(ext(3)-ext(2))*j/(Ny-1.) );
             complex<double> z=0.0;
             for (int itt=0; itt<max_steps; itt++){
                 if (norm(z)>4.) { data(j,i)=itt; break; }
                 z = z*z + z0;     // if |z|>2 the point is not part of mandelbrot set
             }
             if (norm(z)<4.) data(j,i)=max_steps;
         }
     }
     """

def Mandelbrot(ext):
    data = zeros((Nx,Ny),dtype=float)
    for i in range(Nx):
        for j in range(Ny):
            x = ext[0] + (ext[1]-ext[0])*i/(Nx-1.)
            y = ext[2] + (ext[3]-ext[2])*j/(Ny-1.)
            z0 = x+y*1j
            z=0j
            for itr in range(max_steps):
                if abs(z)>2:
                    data[j,i] = itr
                    break
                z = z*z+z0
            if abs(z)<2:
                data[j,i]=max_steps
    return data


# We create class, which inherits from Rectangle, so that it can be called with an Axes
# instance, causing the rectangle to update its shape to match the
# bounds of the Axes
class UpdatingRect(Rectangle):
    def __call__(self, ax):
        self.set_bounds(*ax.viewLim.bounds)  # resize the picture to user specified size
        ax.figure.canvas.draw_idle()         # redraws the picture
        
def ax_update(ax):  # actual plotting routine
    ax.set_autoscale_on(False) # Otherwise, infinite loop

    #Get the number of points from the number of pixels in the window
    dims = ax.axesPatch.get_window_extent().bounds
    width = int(dims[2] + 0.5)
    height = int(dims[2] + 0.5)

    #Get the range for the new area
    xstart,ystart,xdelta,ydelta = ax.viewLim.bounds
    xend = xstart + xdelta
    yend = ystart + ydelta

    ext=array([xstart,xend,ystart,yend])
    #data = Mandelbrot(ext)
    data=zeros((Nx,Ny))
    weave.inline(code, ['data', 'Nx', 'Ny', 'max_steps', 'ext'], type_converters=weave.converters.blitz, compiler = 'gcc')
    
    # Update the image object with our new data and extent
    im = ax.images[-1]  # take the latest object
    im.set_data(data)   # update it with new data
    
    im.set_extent(ext)           # change the extent
    ax.figure.canvas.draw_idle() # finally redraw
                                                                                    

if __name__ == '__main__':

    Nx=300
    Ny=300
    max_steps=50

    ext=array([-2,1,-1,1])
    #data = Mandelbrot(ext)
    data=zeros((Nx,Ny))
    weave.inline(code, ['data', 'Nx', 'Ny', 'max_steps', 'ext'], type_converters=weave.converters.blitz, compiler = 'gcc')
    
    fig,ax=subplots(1,1)
    ax.imshow(data, extent=ext,aspect='equal',origin='lower')
    
    rect = UpdatingRect([0, 0], 0, 0, facecolor='None', edgecolor='black')  # (arguments: xy, width, height, angle=0.0, **kwargs)
    rect.set_bounds(*ax.viewLim.bounds)
    # Connect for changing the view limits
    ax.callbacks.connect('xlim_changed', rect)
    ax.callbacks.connect('ylim_changed', rect)
    ax.callbacks.connect('xlim_changed', ax_update)
    ax.callbacks.connect('ylim_changed', ax_update)
    show()