#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Incrementally render a Mandelbrot set by recursively subdividing intervals.

The idea is that by using interval arithmetic, we can avoid the
majority of the computation involved in rendering the Mandelbrot set,
and thus get a massive speedup.  However, I think it's fair to say
that this didn't actually work out in practice.  I was able to get a
speedup, but not much of one.  Like a third or some shit.

In part this is because the interval-arithmetic slowdown is about
100×, which is a lot: the interval-arithmetic evaluator
interval.mandelbrot_escape_time is able to calculate about 75
Mandelbrot escape times per second on a machine that can do about 7400
per second on native Python complex numbers with mandelbrot_raw from
this file.

Also I was hoping I could use interval arithmetic to calculate a
conservative approximation to the Mandelbrot set, which does, in a
way, work, but it's far, far too conservative, making a huge flood
around the Mandelbrot lake.

"""
import heapq
import random
import time

import pygame

import interval

pix_count = 0

def sub_rect(screen, color, rect):
    (x, y, w, h) = rect
    screen.fill(color, rect)

    global pix_count
    pix_count += w*h
    if pix_count > 512:
        pygame.display.flip()
        while True:
            event = pygame.event.poll()
            if event.type in (pygame.QUIT, pygame.MOUSEBUTTONDOWN):
                raise KeyboardInterrupt
            if event.type == pygame.NOEVENT:
                break
        pix_count = 0
            
def rect(screen, color, rect):
    x, y, w, h = rect
    sub_rect(screen, color, (x, y, w, h))
    sub_rect(screen, color, (x, screen.get_height() - y - h, w, h))

def mandelbrot_raw(c, maxiter):
    z = c
    for m in range(maxiter):
        if abs(z) > 2:
            return m
        z = z * z + c

    return maxiter

conservative = False # slows rendering down by 20× and gives crappy output

def prioritize(prio):
    return prio * random.random()

def render(screen, maxiter, colors, to_screen,
           (outer_iterations, zz), postpone):
    # print screen_rect, zz
    x, y = to_screen((zz.re.min, zz.im.min))
    x, y = int(round(x)), int(round(y))
    x1, y1 = to_screen((zz.re.max, zz.im.max))
    x1, y1 = int(round(x1)), int(round(y1))
    w, h = x1 - x, y1 - y
    if w < 1 or h < 1:
        return

    single_pixel = (w == 1 and h == 1)
    re, im = zz

    iterations = interval.mandelbrot_escape_time(zz, maxiter)

    if single_pixel and (not conservative
                         or maxiter not in iterations):
        zz = interval.ComplexInterval(interval.interval((re.min+re.max)/2),
                                      interval.interval((im.min+im.max)/2))
        re, im = zz
        # iterations = interval.mandelbrot_escape_time(zz, maxiter)
        iterations = interval.interval(
            mandelbrot_raw(re.min + 1j * im.min, maxiter))

    prio = iterations.span() * iterations.max / w / h

    if iterations.max == iterations.min or single_pixel:
        rect(screen, colors[iterations.max], (x, y, w, h))
    elif w*h < 1024 and (not conservative
                         or maxiter not in iterations):
        # Go pixel by pixel now that we're in a small region.  This
        # runs about 10× faster than the whole interval thing, but
        # about ⅓ faster since we're only doing it on regions that
        # might be interesting.
        sub_rect(screen, (2**24-255-255*2**16) | colors[iterations.max],
                 (x, y, w, h))
        pix_width = re.span() / w
        pix_height = im.span() / h
        for jj in range(h):
            im_mid = 1j * (im.min + pix_height * (jj + 0.5))
            for ii in range(w):
                re_mid = re.min + pix_width * (ii + 0.5)
                iterations = mandelbrot_raw(re_mid + im_mid, maxiter)
                rect(screen, colors[iterations], (x+ii, y+jj, 1, 1))
    else:
        sub_rect(screen, (2**24-255) | colors[iterations.max],
                 (x, y, w, h))
        re_span, im_span = re.span(), im.span()
        if w >= h:
            part1 = re.min + re_span / 3
            part2 = re.max - re_span / 3
            # print "a"
            int1 = interval.Interval(re.min, part1)
            postpone(prioritize(prio),
                     (iterations, interval.ComplexInterval(int1, im)))
            # print "b"
            int2 = interval.Interval(part1, part2)
            postpone(prioritize(prio),
                     (iterations, interval.ComplexInterval(int2, im)))
            # print "c"
            int3 = interval.Interval(part2, re.max)
            postpone(prioritize(prio),
                     (iterations, interval.ComplexInterval(int3, im)))
        else:
            part1 = im.min + im_span / 3
            part2 = im.max - im_span / 3
            # print "d"
            int1 = interval.Interval(im.min, part1)
            postpone(prioritize(prio),
                     (iterations, interval.ComplexInterval(re, int1)))
            # print "e"
            int2 = interval.Interval(part1, part2)
            postpone(prioritize(prio),
                     (iterations, interval.ComplexInterval(re, int2)))
            # print "f"
            int3 = interval.Interval(part2, im.max)
            postpone(prioritize(prio),
                     (iterations, interval.ComplexInterval(re, int3)))
    # print "returning"

def main(maxiter=63):
    screen = pygame.display.set_mode((0, 0), pygame.FULLSCREEN)
    w = screen.get_width()
    h = screen.get_height()
    pixels_per_number = h/4.0
    left = -(w/2) / (pixels_per_number)
    top = -2
    to_screen = lambda (re, im): ((re - left) * pixels_per_number,
                                  (im - top) * pixels_per_number)

    gray = (1 << 16) | (1 << 8) | (1 << 0)
    colors = [gray * random.randrange(256) for ii in range(maxiter)]
    colors.append(255)          # blue for maxiter

    start_time = time.time()
    heap = [(0, 
             (interval.Interval(0, maxiter),
              interval.ComplexInterval(interval.Interval(left, -left),
                                       interval.Interval(-2., 0.))))]
    while heap:
        item = heapq.heappop(heap)[1]
        render(screen,
               maxiter,
               colors,
               to_screen,
               item,
               lambda prio, item: heapq.heappush(heap, (prio, item)))

    pygame.display.flip()

    duration = time.time() - start_time
    print "%ss to render %s pixels, %s pixels/s" % (duration, w*h, w*h/duration)

    while pygame.event.poll().type not in (pygame.QUIT, pygame.MOUSEBUTTONDOWN):
        pygame.time.delay(100)

if __name__ == '__main__':
    main()