import pygame, colorsys
import numpy as np
from numba import jit

# Set the width and height of the window
width, height = 1100, 800

# Define the region of the complex plane to visualize
xmin, xmax = -2, 1
ymin, ymax = -1, 1

# Set the maximum number of iterations for the Mandelbrot calculation
max_iterations = 100

# Initialize Pygame
pygame.init()

# Create a window
screen = pygame.display.set_mode((width, height))
pygame.display.set_caption("Mandelbrot Set")

@jit(nopython=True)
def mandelbrot(c):
    z = 0
    for i in range(max_iterations):
        z = z**2 + c
        if abs(z) > 2:
            return i
    return max_iterations

# Function to convert iterations to a smooth gradient in the HSV color space
def smooth_color(iterations):
    if iterations == max_iterations:
        return 0, 0, 0  # Black for points in the set
    hue = 0.7 + iterations / max_iterations
    saturation = 1.0
    value = iterations / max_iterations
    return [int(c * 255) for c in colorsys.hsv_to_rgb(hue, saturation, value)]

# Function to update the display
def update_display():
    for x in range(width):
        for y in range(height):
            c = complex(xmin + (x / width) * (xmax - xmin), ymin + (y / height) * (ymax - ymin))
            iterations = mandelbrot(c)
            color = smooth_color(iterations)
            screen.set_at((x, y), color)
    pygame.display.flip()

# Run the Pygame event loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    update_display()

# Quit Pygame
pygame.quit()