from math import acos, sin, cos, pi, sqrt
from pyx import *
import os

def draw_layer(color):
   global c
  for square in layer:
     if square[1] > .0001:
      c.stroke(rect, [ deco.filled([color]), color,
         trafo.translate(square[0][0],square[0][1]) *
         trafo.rotate(square[2]*180/pi) *
         trafo.scale(square[1]) ])

def next_left():
  return set( [ ((t[0]-s*sin(r), t[1]+s*cos(r)), s*ls, r+theta )
    for (t,s,r) in layer ] )

def next_right():
     return set( [ ((t[0]+s*(ls*cos(r+theta)-sin(r)),
                       t[1]+s*(cos(r)+ls*sin(r+theta))),
                       s*rs, r-phi )
                   for (t,s,r) in layer ] )


if __name__ == "__main__":

     ## INITIALIZE
     N = 15                  # NUMBER OF LAYERS
     imgname = 'Ptree_7-7'   # IMAGE NAME
     ls = .7                 # LEFT BRANCH SIZE
     rs = .7                 # RIGHT BRANCH SIZE ; ls+rs >= 1

     theta = acos( (ls**2-rs**2+1) / (2*ls) )
     phi = acos( (rs**2-ls**2+1) / (2*rs) )
     layer = set([ ((0.0,0.0), 1.0, 0.0) ])
     c = canvas.canvas()
     rect = path.rect(0,0,1,1)

     for i in xrange(N-1):
           draw_layer(color.cmyk.RawSienna)
           layer = next_left() | next_right()
           print 'Layer', i+1

     draw_layer(color.rgb.green)
     c.writePDFfile(imgname)