# fractal_tennis.py  -- by leonardo maffi, V.1.1, July 26 2006.

from random import random
import pygame # http://www.pygame.org/download.shtml

sierpinsky_triangle = [[0.5, 0.0, 0.0, 0.5, 0.0, 0.0, 0.33], # The last number is the probability
                       [0.5, 0.0, 0.0, 0.5, 0.5, 0.0, 0.66],
                       [0.5, 0.0, 0.0, 0.5, 0.25, 0.5, 1.0]]

fern_leaf = [[0.844, 0.0255, -0.0255, 0.849, 0.0, 3.0, 0.85],
             [-0.155, 0.235, 0.195, 0.186, 0.0, 1.2, 0.92],
             [0.155, -0.235, 0.195, 0.186, 0.0, 3.0, 0.992],
             [0.0, 0.0, 0.0, 0.17, 0.0, 0.0, 1.0]]


def transformation_apply(px, py, t):
    return px*t[0] + py*t[1] + t[4], px*t[2] + py*t[3] + t[5]


def transformation_normalize(t, nx, ny, ntests=10000):
    minx = 1e5; maxx = -1e5
    miny = 1e5; maxy = -1e5
    px, py = 0.0, 0.0
    for i in xrange(ntests):
        r = random()
        n = 0
        while r > t[n][6]:
            n += 1
        px, py = transformation_apply(px, py, t[n])
        minx = min(px, minx)
        maxx = max(px, maxx)
        miny = min(py, miny)
        maxy = max(py, maxy)
    mx = nx / (maxx-minx)
    my = ny / (maxy-miny)
    return minx, miny, mx, my


def tennis_plot(screen, nx, ny, transf, npoints=100000):
    white = (255, 255, 255)
    sx, sy, mx, my = transformation_normalize(transf, nx, ny)
    px, py = 0.0, 0.0
    for i in xrange(npoints):
        r = random()
        for tr in transf:
            if r <= tr[6]:
                break
        px, py = transformation_apply(px, py, tr)
        if i > 10: # Skip first bad points
            screen.set_at( (int((px-sx)*mx), int((py-sy)*my)), white ) # Slow
        if not (i % 10000):
            # refresh screen once in a while
            pygame.display.flip()


def tennis_plot_grey_shades(screen, nx, ny, transf, npoints=100000):
    # This is slower than tennis_plot
    try:
        import Numeric as N # numarray or scipy can be fine too
        import pygame.surfarray as surfarray
    except ImportError:
        raise ImportError, "Numeric and Surfarray are required."
    screen_arr = N.zeros((nx, ny), typecode="l")
    nx1 = nx - 1
    ny1 = ny - 1
    sx, sy, mx, my = transformation_normalize(transf, nx, ny)
    px, py = 0.0, 0.0

    for i in xrange(npoints):
        r = random()
        for tr in transf:
            if r <= tr[6]:
                break
        px, py = transformation_apply(px, py, tr)
        if i > 10: # Skip first bad points
            x, y = int((px-sx)*mx), int((py-sy)*my)
            x = max(0, min(nx1, x))
            y = max(0, min(ny1, y))
            screen_arr[x, y] += 1
        if not (i % 20000):
            # refresh screen once in a while.
            # Here it can be called a function that does a good normalizing,
            #   taking 5% to 95% of values, working with RGB too.
            scr_arr_scaled = screen_arr * 10
            N.putmask(scr_arr_scaled, scr_arr_scaled>255, 255)
            surfarray.blit_array(screen, scr_arr_scaled)
            pygame.display.flip()


def pygame_plot(do_plot, nx=640, ny=480, **args):
    pygame.init()
    screen = pygame.display.set_mode((nx, ny), 0, 8) #greyscale screen!
    pygame.display.set_palette( [(i,i,i) for i in xrange(256)] ) #greyscale screen!
    pygame.display.set_caption(do_plot.__name__)
    do_plot(screen, nx, ny, **args)
    pygame.display.flip()
    while True:
        e = pygame.event.wait()
        if e.type == pygame.QUIT:
            break


try: # Use psyco if possible
    import psyco
except:
    pass
else:
    psyco.bind(tennis_plot)
    psyco.bind(tennis_plot_grey_shades)

#pygame_plot(tennis_plot, nx=500, ny=700, transf=fern_leaf, npoints=200000)
#pygame_plot(tennis_plot, nx=500, ny=500, transf=sierpinsky_triangle, npoints=100000)
#pygame_plot(tennis_plot_grey_shades, nx=500, ny=500, transf=sierpinsky_triangle, npoints=400000)
pygame_plot(tennis_plot_grey_shades, nx=500, ny=700, transf=fern_leaf, npoints=400000)