"""
num_pyramid.py -- V.1.0, Dec 31 2006 by leonardo maffi

Requires the CSP module, see http://www.fantascienza.net/leonardo/so/index.html
Psyco optional.

The code is quite slow, it requires about 40 seconds on a PIII 500 MHz with Python 2.5 + Psyco.
Using a system of equations it can probably be solverd in microseconds.
"""

from CSP import Problem, timing

try: # Use Psyco if available
    import psyco
    psyco.full()
except ImportError:
    pass


def pyramid_of_numbers():
    print """Pyramid of numbers problem:
    Each brick of the pyramid is the sum of the two bricks situated below this brick.
    For the tree missing numbers at the base of the pyramid: the one of the middle is
    the sum of the two other. How can we fill the pyramid of numbers?
    See also: http://xunor.free.fr/en/riddles/auto/pyramidnb.php
    """

    pyramid = """
              151
            ??   ??
         40   ??   ??
       ??   ??   ??   ??
    ??   11   ??    4   ??"""

    print "Problem:", pyramid, "\n"

    data0 = [row.split() for row in pyramid.splitlines() if row.strip()]
    data = [[(None if x=="??" else int(x)) for x in row] for row in data0]
    print "Encoded input data:\n", data, "\n"

    def pos(row, col):
        # To convert a position to number.
        # This trick allows you to use integers as variables instead of tuples,
        #   speeding up the computation. Using pos(r,c) you keep the ability
        #   to manage pairs in a simple way.
        return row*5 + col

    def ppos(row, col): # to print a position
        return "<%d,%d>" % (row, col)

    p = Problem("recursivebacktracking")
    print "Solver type: Recursive Backtracking"

    print "Add variables:"
    domain = range(0, 152) # hypotetical
    for nrow, row in enumerate(data):
        for ncol, el in enumerate(row):
            p.var(pos(nrow, ncol), domain)
            print ppos(nrow, ncol),
        print
    print

    print "Add known values:"
    for nrow, row in enumerate(data):
        for ncol, el in enumerate(row):
            if el is not None:
                p.rule(lambda var, val=el: var == val, [pos(nrow, ncol)])
                print ppos(nrow,ncol), "==", el
    print

    print "Each brick of the pyramid is the sum of the two bricks situated below this brick:"
    for nrow in xrange(4):
        for ncol in xrange(nrow+1):
            p.rule(lambda v1, v2, v3: v1 == v2 + v3, [pos(nrow,ncol), pos(nrow+1,ncol), pos(nrow+1,ncol+1)])
            print ppos(nrow,ncol), "==", ppos(nrow+1,ncol), "+", ppos(nrow+1,ncol+1)
    print

    print "In the last row the one in the middle is the sum of the two others:"
    p.rule(lambda v1, v2, v3: v1 == v2 + v3, [pos(4,2), pos(4,0), pos(4,4)])
    print ppos(4,2), "==", ppos(4,0), "+", ppos(4,4)
    print

    sol, t = timing(p.solution)
    # About 40 seconds on a PIII 500 MHz with Python 2.5 + Psyco

    if sol:
        print "Solution (%s s):" % str(round(t, 2))
        for nrow, row in enumerate(data):
            for ncol in xrange(len(row)):
                print sol[pos(nrow, ncol)],
            print
    else:
        print "Solution not found (%s s)." % str(round(t, 2))

    solution = """\
              151
            81   70
         40   41   29
      16   24   17   12
    5   11   13   4     8"""


pyramid_of_numbers()