#!/usr/bin/env python
import psyco
psyco.log()
psyco.full()
import pygame
from pygame.locals import *
from math import log
import time

def _default_gen():
    blank = [0 for i in xrange(256)]
    rising = [i for i in xrange(256)]
    full = [255 for i in xrange(256)]
    return zip(rising + full + full, blank + rising + full, blank + blank + rising)*10
class Mandelbrot(object):
    
    def __init__(self, surface, antialiasing=1, smooth=True, max_iteration=768, minimum_real=-2.0, maximum_real=1.0, minimum_imaginary=-1.2, escape=2, degree=2, gen_colourtable=_default_gen):
        self.colourtable = gen_colourtable()
        self.degree = degree
        self.smooth = smooth
        self.escape = escape
        self.escape_squared = escape*escape
        self.surface = surface
        self.max_iteration = max_iteration
        self.height = surface.get_height()*antialiasing
        self.width = surface.get_width()*antialiasing
        if antialiasing < 1:
            raise ValueError, "antialiasing must be 1 or greater"
        else:
            self.antialiasing = antialiasing
        
        self.minimum_real = minimum_real
        self.maximum_real = maximum_real
        self.minimum_imaginary = minimum_imaginary
        self.maximum_imaginary = self.minimum_imaginary+(self.maximum_real-self.minimum_real)*self.height/self.width
        self.x_size_factor = (self.maximum_real-self.minimum_real)/(self.width-1);
        self.y_size_factor = (self.maximum_imaginary-self.minimum_imaginary)/(self.height-1);
        
        self.rendered = False
        
    def render(self):
        surface = self.surface
        for x in xrange(self.width):
            x0 = self.minimum_real + x*self.x_size_factor
            for y in xrange(self.height):
                y0 = self.maximum_imaginary - y*self.y_size_factor
                
                c = x0 + y0*1j
                z = 0+0j
                iteration = 0
                
                for i in xrange(self.max_iteration):
                    #xtemp = x*x - y*y + x0
                    #ytemp = 2*x*y + y0
                    
                    #x = xtemp
                    #y = ytemp
                    
                    iteration = i
                    #(a+bi)(c+di) = a*c + c*bi + a*di+bi*di, just FTR
                    z = z**self.degree + c
                    if z.real*z.real + z.imag*z.imag > self.escape_squared: #possibly wrong?
                        break
                colour = self.gen_colour_index(iteration, z)
                
                self.plot((x,y),colour)
        #self.surface.blit(pygame.transform.rotozoom(self.surface, 0, 1.0/self.antialiasing), (0,0))
        self.rendered = True
    def gen_colour_index(self, iteration, z):
        if iteration == self.max_iteration - 1:
            return 0
        else:
            if self.smooth == True:
                return int((iteration + (log(2*log(self.escape))-log(log(abs(z))))/log(self.degree)) * 10)
            else:
                return iteration*10
    
    def plot(self, (x,y), colour):
        self.surface.set_at((x, y), self.colourtable[colour])
        #self.surface[x][y] = self.colourtable[colour]

class Cursor(object):
    def __init__(self, display, size=5):
        self._base_w = display.dimensions[0] / 20
        self._base_h = display.dimensions[1] / 20
        self._display = display
        self._bgcolor = 0,0,0
        
        self.size = size
    
    def reset_image(self):
        self.img = pygame.Surface((self.width, self.height))
        pygame.draw.rect(self.img, (255,255,255), self.img.get_rect(), 4)
        
    
    def _get_size(self):
        return self._size
    
    def _set_size(self, size):
        if 0<=size<=20:
            self._size = size
            self.width = self._base_w * size
            self.height = self._base_h * size
            self.reset_image()
        else:
            raise ValueError
    size = property(_get_size, _set_size)

class Display(object):
    def __init__(self):
        pygame.init()
        flags = 0
        self.dimensions = (100, 100)
        self.screen = pygame.display.set_mode(self.dimensions, flags)
        pygame.display.set_caption("Mandelbrot test")
        self.cursor = Cursor(self)
    
    def run(self):
        mandelsheet = pygame.Surface(self.dimensions)
        cursorsheet = pygame.Surface(self.dimensions)
        cursorsheet.set_colorkey((0,0,0))
        self.mandelbrot = Mandelbrot(mandelsheet)
        while True:
            for event in pygame.event.get():
                if event.type == MOUSEBUTTONDOWN and event.button == 1:
                    if self.mandelbrot.rendered:
                        x, y = event.pos
                        
                        #x0 = (x - self.mandelbrot.minimum_real) / self.mandelbrot.x_size_factor
                        #x1 = (x + self.cursor.width - self.mandelbrot.minimum_real) / self.mandelbrot.x_size_factor
                        #y0 = (-y + self.mandelbrot.maximum_imaginary)/self.mandelbrot.y_size_factor
                        x0 = self.mandelbrot.minimum_real + x*self.mandelbrot.x_size_factor
                        x1 = self.mandelbrot.minimum_real + (x+self.cursor.width)*self.mandelbrot.x_size_factor
                        y0 = self.mandelbrot.maximum_imaginary - (y + self.cursor.height)*self.mandelbrot.y_size_factor
                        
                        minimum_real=x0
                        maximum_real=x1
                        minimum_imaginary=y0
                        
                        self.mandelbrot = Mandelbrot(mandelsheet, minimum_real=minimum_real, maximum_real=maximum_real, minimum_imaginary=minimum_imaginary)
                elif event.type == MOUSEBUTTONDOWN and event.button == 4:
                    #wheel up
                    try:
                        self.cursor.size -= 1
                    except ValueError:
                        pass
                elif event.type == MOUSEBUTTONDOWN and event.button == 5:
                    #wheel down
                    try:
                        self.cursor.size += 1
                    except ValueError:
                        pass
                elif event.type == QUIT:
                    return
                elif event.type == KEYDOWN and event.key == K_ESCAPE:
                    return
            if not self.mandelbrot.rendered:
                self.mandelbrot.render()
                
            self.screen.blit(mandelsheet, (0,0))
            cursorsheet.fill((0,0,0))
            cursorsheet.blit(self.cursor.img, pygame.mouse.get_pos())
            self.screen.blit(cursorsheet, (0,0))
            pygame.display.flip()
            #return #for timing purposes, REMOVE WHEN DONE OPTIMIZING
            #raw_input("Done. View, and press <enter> to exit. ")
    
if __name__ == "__main__":
    disp = Display()
    #s_time = time.time()
    disp.run()
    #print "time: %ss" %(time.time() - s_time)