Chaos game en python

chaos-game.py

import sys
import math
import random


def midpoint(p0, p1):
    (x0, y0)= p0
    (x1, y1)= p1
    return ((x0 + x1)/2, (y0 + y1)/2)

def rotate(center, a, p):
    (x0, y0)= center
    (x, y)= p
    x-= x0
    y-= y0
    return (x0+math.cos(a)*x - math.sin(a)*y,
            y0+math.sin(a)*x + math.cos(a)*y)

def polygon(center, p, n):
    res=[]
    for i in range(math.ceil(n)):
        res.append(p)
        p= rotate(center, 2*math.pi/n, p)
    return res

def brighten(screen, p, white):
    (x,y)= p
    x= int(x)
    y= int(y)
    screen[y][x]= min(screen[y][x]+1, white)

def create_screen(init, w, h):
    return [[init for x in range(w)] for y in range(h)]

def print_image(white, screen):
    print("P2")
    print(len(screen[0]), len(screen[1]), sep=" ")
    print(white)
    for line in screen:
        for c in line :
            print(c, end=" ")
    print()

def main(argv):
    n_edges= float(argv[0]) if len(argv)>0 else 3.
    w= 1024
    h= 1024
    black= 0
    white= 255
    n= int(n_edges * ( int(argv[1] if len(argv) > 1 else 10000000)))
    poly= polygon((w/2, h/2), (w/2, 1.75*h/2), n_edges)
    screen= create_screen(black, w, h)
    prev_idx= 0
    p= poly[prev_idx]
    for i in range(n):
      idx= random.randrange(len(poly))
      while idx == prev_idx:
        idx= random.randrange(len(poly))
      prev_idx= idx
      p= midpoint(poly[idx], p)
      brighten(screen, p, white)
    print_image(white, screen)
if __name__ == "__main__":
    main(sys.argv[1:])