import numpy as np
import matplotlib.pyplot as plt

'''
z₀ = 0
zₙ₊₁ = zₙ² + c
'''
def is_mandelbrot(complex_num, max_iter=100):
    z = 0
    for n in range(max_iter):
        z = z**2 + complex_num
        if abs(z) > 2:
            return False

    return True

def generate_mandelbrot_points(resolution=100):
    real = np.linspace(-2, 0.8, resolution)
    imag = np.linspace(-1.4, 1.4, resolution)
    points = []

    for r in real:
        for i in imag:
            c = complex(r, i)
            if is_mandelbrot(c):
                points.append(c)

    return points


def plot_mandelbrot(points):
    real_parts = [p.real for p in points]
    imag_parts = [p.imag for p in points]
    
    plt.figure(figsize=(10,10))
    plt.scatter(real_parts, imag_parts,s=1,c="black")
    plt.axis("equal")
    plt.show()

if __name__ == "__main__":
    mandelbrot_points = generate_mandelbrot_points(resolution=100)
    print(len(mandelbrot_points), "points in the set")

    plot_mandelbrot(mandelbrot_points)