import math
import matplotlib.pyplot as plt
import numpy as np

import time


def mandel_iter(c,max_escape_time):
    z = c
    for count in range(1,max_escape_time+1):
        if abs(z) > 2:
            return count
        z = z*z + c

    return max_escape_time


def compute_mandelbrot(xmin,xmax,ymin,ymax,Nx,Ny,max_escape_time=1000,
                        plot_filename=None,cur_cmap=plt.cm.prism,totime=False):


    t0 = time.clock()


    X = np.linspace(xmin, xmax, Nx, endpoint=True)  #real values
    Y = np.linspace(ymin,  ymax, Ny, endpoint=True) #imag values
    Y = np.flipud(Y)    #reverse Y to get the right direction on the plot

    #A little bit faster with list instead of array for some reason
    res = [[0 for x in range(len(X))] for y in range(len(Y))]
    # res = np.empty(shape=(Ny,Nx),dtype=np.float)

    # do computation
    for y_index,imag in enumerate(Y):
        for x_index,real in enumerate(X):
            res[y_index][x_index] = mandel_iter(complex(real,imag),max_escape_time)

    t1 = time.clock()


    t1 = time.clock()

    if totime:
        print("Computation time (basic python) : " + str(t1-t0))

    #Draw if plot_filename provided
    if plot_filename:
        #Mask for drawing, to get points in mandelbrot black
        Z = np.asarray(res,dtype=np.float64)
        Z[Z == max_escape_time] = np.NaN
        masked_array = np.ma.array(Z, mask=np.isnan(Z))

        cur_cmap.set_bad('black',1.)
        plt.imshow(Z, cmap = cur_cmap, extent =(xmin, xmax, ymin, ymax))
        plt.savefig(plot_filename+'.png')
        plt.show()

    res = np.asarray(res,dtype=np.float64)
    return res #return non-masked results, with max_escape_time instead of np.NaN

if __name__ == "__main__":
    compute_mandelbrot(
        xmin = -2,
        xmax =  1,
        ymin=-1,
        ymax=1,
        Nx=300,
        Ny=300,
        max_escape_time=1000,
        # plot_filename="m1",
        plot_filename=None,
        totime=True
        )