# mandel.py
# https://www.geeksforgeeks.org/mandelbrot-fractal-set-visualization-in-python/
# https://users.math.yale.edu/public_html/People/frame/Fractals/MandelSet/welcome.html


import math
import matplotlib.pyplot as plt


MAX_ITER = 50
Z_MAX = 2

# parameters for the plot
imWidth, imHeight = 500, 500
xmin, xmax = -2, 1.0
xwidth = xmax - xmin
ymin, ymax = -1.5, 1.5
yheight = ymax - ymin

# mandel[y][x]
mandel = [[0 for x in range(imWidth)] for y in range(imHeight)]

for x in range(imWidth):
  for y in range(imHeight):
    # Map pixel position to a point in the complex plane
    c = complex((x / imWidth * xwidth) + xmin,
                (y / imHeight * yheight) + ymin)
    z = 0
    for i in range(MAX_ITER):
      if abs(z) >= Z_MAX:
        mandel[y][x] = i
        break
      else: 
        z = z**2 + c

plt.imshow(mandel, interpolation='nearest', cmap="hot")
plt.xticks([])
plt.yticks([])
plt.title("Mandelbrot Set")
plt.show()