import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

from mandelbrot import mandelbrot


def animategif():
    x_start, y_start = -2, -1.5  # an interesting region starts here
    width, height = 3, 3  # for 3 units up and right
    density_per_unit = 100  # how many pixels per unit

    # real and imaginary axis
    re = np.linspace(x_start, x_start + width, width * density_per_unit)
    im = np.linspace(y_start, y_start + height, height * density_per_unit)

    fig = plt.figure(figsize=(5, 5))  # instantiate a figure to draw
    ax = plt.axes()  # create an axes object

    def animate(i):
        ax.clear()  # clear axes object
        ax.set_xticks([], [])  # clear x-axis ticks
        ax.set_yticks([], [])  # clear y-axis ticks

        X = np.empty((len(re), len(im)))  # re-initialize the array-like image
        threshold = round(1.15 ** (i + 1))  # calculate the current threshold

        # iterations for the current threshold
        for i in range(len(re)):
            for j in range(len(im)):
                X[i, j] = mandelbrot(re[i], im[j], threshold)

        # associate colors to the iterations with an iterpolation
        img = ax.imshow(X.T, interpolation="bicubic", cmap='magma')

        return [img]

    anim = animation.FuncAnimation(fig, animate, frames=45, interval=120, blit=True)
    return anim