from PIL import Image, ImageDraw
import gradio as gr


def mandelbrot(c, MAX_ITER):
    z = 0
    n = 0
    while abs(z) <= 2 and n < MAX_ITER:
        z = z*z + c
        n += 1
    return n


def mandelbrot_set(MAX_ITER):
    # Image size (pixels)
    WIDTH = 600
    HEIGHT = 400

    # Plot window
    RE_START = -2
    RE_END = 1
    IM_START = -1
    IM_END = 1

    palette = []

    im = Image.new('RGB', (WIDTH, HEIGHT), (0, 0, 0))
    draw = ImageDraw.Draw(im)

    for x in range(0, WIDTH):
        for y in range(0, HEIGHT):
            # Convert pixel coordinate to complex number
            c = complex(RE_START + (x / WIDTH) * (RE_END - RE_START),
                        IM_START + (y / HEIGHT) * (IM_END - IM_START))
            # Compute the number of iterations
            m = mandelbrot(c, MAX_ITER)
            # The color depends on the number of iterations
            color = 255 - int(m * 255 / MAX_ITER)
            # Plot the point
            draw.point([x, y], (color, color, color))

    return im



iface = gr.Interface(mandelbrot_set, 
                    gr.inputs.Slider(1, 100, 1), 
                    "pil",
                    title='Mandelbrot Set',
                    description='Mandelbrot Set generation with n iterations <img src="https://upload.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/767px-Cat_November_2010-1a.jpg" width=100>')

iface.launch()