import random
import math

# Boltzmann generator for binary trees with the specification T = E + ZT^2
# and generator for the pointed class.  This improves the performance

def BoltzmannTree(x):
    # set up	
    u = random.random()

    # handle union		
    if u < 1/((1-math.sqrt(1-4*x))/(2*x)):
        return 'E'
    
    # handle product
    return ['Z', BoltzmannTree(x), BoltzmannTree(x)]

def BoltzmannPointedTree(x):
    # set up
    # in the maple case I was lazy and did the differentiation and substitution
    # in the file.  Here I've done it in advance since plain python doesn't
    # have symbolic features.
    T = (1-math.sqrt(1-4*x))/(2*x)
    pointedT = (1-2*x-math.sqrt(1-4*x))/(2*math.sqrt(1-4*x)*x*x)
    u = random.random()

    # handle union
    if u < x*T/pointedT:
        return ['Z', BoltzmannTree(x), BoltzmannTree(x)]
    if u < x*T/pointedT + x*pointedT*T/pointedT:
        return ['Z', BoltzmannPointedTree(x), BoltzmannTree(x)]
    
    return ['Z', BoltzmannTree(x), BoltzmannPointedTree(x)]


print BoltzmannPointedTree(0.24)
print BoltzmannPointedTree(0.2499)
