from scipy import *
from pylab import *
import time
import weave
import imanc

code="""
   #line 8 "manp3.py"
   using namespace std;
   for (int i=0; i<Nx; i++){
      for (int j=0; j<Ny; j++){
          double x = ext(0) + (ext(1)-ext(0))*i/(Nx-1.);
          double y = ext(2) + (ext(3)-ext(2))*j/(Ny-1.);
          complex<double> z0(x,y);
          complex<double> z=0;
          for (int itr=0; itr<max_steps; itr++){
             if (norm(z)>4.){
                data(j,i) = itr;
                break;
             }
             z = z*z + z0;
          }
          if (norm(z)<4.) data(j,i) = max_steps;
      }
   }
"""

def Mand(ext, max_steps):
    data = zeros( (Nx,Ny) )
    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
def Mandc(ext, max_steps):
    data = zeros((Ny,Nx));
    ext=array(ext)
    weave.inline(code, ['data', 'Nx', 'Ny', 'max_steps', 'ext'], type_converters=weave.converters.blitz, compiler = 'gcc')
    return data
def Mandc2(ext, max_steps):
    data = zeros((Ny,Nx));
    imanc.mand(data, Nx, Ny, max_steps, ext)
    return data
    
if __name__ == '__main__':

    Nx = 1000
    Ny = 1000
    max_steps = 50

    ext = [-2,1,-1,1]
    t0 = time.clock()
    data = Mandc2(ext, max_steps)
    print 'time=', time.clock()-t0
    
    imshow(data, extent=ext)
    show()