import pygame
import numpy
from pygame.locals import *
from sys import exit
from random import randint

width, height = 800, 800

pygame.init
screen = pygame.display.set_mode((width, height), 0, 32)

# Currently draws the Sierpinski triangle. Change these to draw
# different shapes. What this does precisely is to transform the
# point (x, y) into (ax + by + c, dx + ey + f)
# for a = a1, a2, a3, b = b1, b2, b3, etc.
a1, b1, c1, d1, e1, f1 = 0.5, 0, 0, 0.5, 0, 346.41
a2, b2, c2, d2, e2, f2 = 0.5, 0, 0, 0.5, 400, 346.41
a3, b3, c3, d3, e3, f3 = 0.5, 0, 0, 0.5, 200, 0

def init_array():
    a = numpy.zeros([height, width], int)
    return a

def next_iteration(a):
    s = numpy.zeros([height, width], int)
    for i in range(width):
        for j in range(height):
            if a[j][i] == 1:
                x1, y1 = a1*i + b1*j + e1, c1*i + d1*j + f1
                x2, y2 = a2*i + b2*j + e2, c2*i + d2*j + f2
                x3, y3 = a3*i + b3*j + e3, c3*i + d3*j + f3
				
                if(x1 < width and y1 < height
                    and x2 < width and y2 < height
                    and x3 < width and y3 < height
                    and x1 > 0 and y1 > 0
                    and x2 > 0 and y2 > 0
                    and x3 > 0 and y3 > 0):
                        s[int(y1)][int(x1)] = 1
                        s[int(y2)][int(x2)] = 1
                        s[int(y3)][int(x3)] = 1
    return s

a = init_array()

screen.lock()
for i in range(width):
    for j in range(height):
        if(a[j][i] == 1):
            screen.set_at((i, j), (255, 255, 255))
screen.unlock()

pygame.display.update()

while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            exit()
        if event.type == MOUSEMOTION or event.type == MOUSEBUTTONDOWN:
            (b1, b2, b3) = pygame.mouse.get_pressed()
            if(b1):
                (x,y) = pygame.mouse.get_pos()
                a[y][x] = 1
                screen.set_at((x, y), (255, 255, 255))
                pygame.display.update()
        if event.type == KEYDOWN:
            if event.key == K_ESCAPE:
                exit()
            else:
                screen.fill((0, 0, 0))
                a = next_iteration(a)
                for i in range(width):
                    for j in range(height):
                        if(a[j][i] == 1):
                            screen.set_at((i, j), (255, 255, 255))
                pygame.display.update()