#!/usr/bin/env python
#Fractalss ver 0.1
#Vadim Kataev 2006
#www.compuvisor.net
#www.technopedia.org
#vkataev at gmail.com
# Modified by Charles Wicksteed, June 2022

import pygame, time
import pygame.locals

SCREEN_WIDTH=120; SCREEN_HEIGHT=200
CENTER_X=SCREEN_WIDTH/2; CENTER_Y=SCREEN_HEIGHT/2

max_iteration = 70
scale = 3.0/(SCREEN_HEIGHT*500.0)

def init():
  global screen
  pygame.init()
  screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
  pygame.display.set_caption('Mandelbrot')

def main():
  screen.fill((255, 255, 255))
  if (draw_field(screen) == "DONE"):
    # wait for user to close window
    running = True
    while running:
      event = pygame.event.wait()
      if event.type == pygame.locals.QUIT:
        running = False
  # else if draw_field() returns "QUIT", then quit immediately

# Mandelbrot Set
def iteration(cx, cy):
  i = 0
  zx = 0
  zy = 0
  while zx * zx + zy * zy < 4.0 and i<max_iteration:
    xtemp = zx * zx - zy * zy + cx
    zy = 2 * zx * zy + cy
    zx = xtemp
    i+=1
  return i

def set_color(i):
  if i == max_iteration:
    col = (0,0,0)
    return col
  else:
    # https://www.pygame.org/docs/ref/color.html
    c = pygame.Color(0, 0, 0)
    # hue saturation lightness alpha
    c.hsla = (int(i * 3.5) % 360, 100, 50, 100)
    return c

def draw_field(scr):
  for scr_x in range(SCREEN_WIDTH):
    for scr_y in range(SCREEN_HEIGHT):
      x = (scr_x - CENTER_X) * scale - 0.001
      y = (scr_y - CENTER_Y) * scale - 0.75
      iter_value = iteration(x, y)
      col = set_color(iter_value)
      pygame.draw.line(scr, col, (scr_x, scr_y), (scr_x, scr_y))
    # after each column, update screen and check for quit
    pygame.display.flip()
    for event in pygame.event.get():
      if event.type == pygame.locals.QUIT:
        return "QUIT"
  # finished whole pattern
  return "DONE"

init()
main()
pygame.quit()