from sys import argv, exit
from random import randrange, choice
from math import log

def depth(token):
    return int(log(token)/log(2))

def next(token, maxDepth):
    if depth(token) == maxDepth-1:
        opts = ["up"]
    elif token == 1:
        opts = ["left", "right"]
    else:
        opts = ["up", "left", "right"]
    do = choice(opts)
    if do == "up":
        return token / 2
    elif do == "left":
        return 2*token
    else:
        return 2*token + 1

def path(token):
    res = []
    while token >= 1:
        res.append(token)
        token /= 2

    return res

if __name__ == "__main__":
    if len(argv) < 2:
        print "Usage: tree [depth]"
        exit()

    maxDepth = int(argv[1])
    nodes = 2**maxDepth #nodes+1
    print "Complete binary tree of depth", maxDepth, ";", nodes-1, "nodes."

    token = randrange(1, nodes)
    print "Placing a token at", token, "."

    samples = 5000
    sAtD = [0.0] * maxDepth
    under = {}
    firstAtHeight = {}
    firstAtHeight[depth(token)] = 1

    for i in range(samples):
        token = next(token, maxDepth)
        sAtD[depth(token)] += 1

        if not firstAtHeight.has_key(depth(token)):
            firstAtHeight[depth(token)] = i

        for ancestor in path(token):
            if under.has_key(ancestor):
                under[ancestor] += 1
            else:
                under[ancestor] = 1

    print "Time at depth:"
    for d in range(maxDepth):
        print d, ":\t", (0.0+sAtD[d])/samples
    print

    ancestorsAt = [0.0] * maxDepth
    for anc,ct in under.items():
        ancestorsAt[depth(anc)] += 1

    print "Ancestors represented at:"
    for d in range(maxDepth):
        print d, ":\t", (0.0+ancestorsAt[d]), "out of", 2**d
    print

    print "First visit times:"
    for d in range(maxDepth):
        if firstAtHeight.has_key(d):
            print "Depth", d, "at time", firstAtHeight[d], "."
        else:
            print "Depth", d, "never."
