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

N = 1000
inital_z = 0

epsilon = 0.01
center_position = [-0.7, 0.46]

def f(z, C):
    return z*z + C

def convergence_judgment(C):
    z = inital_z
    i = 0
    while i < 100:
        z = f(z, C)
        if abs(z) > 2:
            return i
        i += 1
    return i

def show(grid):
    plt.figure(figsize=(15,15))
    plt.axis("off")
    plt.tick_params(labelbottom=False, labelleft=False, labelright=False, labeltop=False)
    plt.imshow(grid, vmin=0, vmax=100, cmap="Blues_r")
    plt.savefig("save/img_x{}_y{}_e{}.png".format(center_position[0], center_position[1], epsilon))
    plt.show()

def main():
    grid = np.empty((N, N), dtype=np.uint8)
    points_x = np.linspace(center_position[0]-epsilon, center_position[0]+epsilon, N)
    points_y = np.linspace(center_position[1]-epsilon, center_position[1]+epsilon, N)

    for n, ReC in tqdm(enumerate(points_x)):
        for m, ImC in enumerate(points_y):
            C = ReC + ImC*1J
            grid[m, n] = convergence_judgment(C)

    show(grid)

if __name__ == "__main__":
    main()