# This is Python 3 code
from zellegraphics import * 
# We use this package in CSSE 120.  FOr installation instructions, 
# see http://www.rose-hulman.edu/class/cs/resources/Python/installation.htm

def optimalBST(n, a, b, R, win, indent, squareSize):
    """ n = number of nodes in tree
        a = frequencies of successful search (internal nodes of the tree)
        b = frequencies of unsuccessful search (external nodes)
        R = position of the root of optimal (sub) tree (calculated by this pworgram)
        All these are described in the OptimalBST PowerPoint slides. """
    print("a = ",a)
    print("b = ", b)
    C = makeTable(n)   # C and W are described in the PowerPoint slides.
    W = makeTable(n)
    # initialize the main diagonal
    for i in range(n + 1):
        R[i][i] = i
        W[i][i] = b[i]
        # Draw this cell of the table in the given window.
        drawSquare(i, i, W[i][i], C[i][i], R[i][i], win, indent, squareSize )
    # Now populate each of the n upper diagonals:
    for d in range(1, n+1): # fill in this diagonal
        # The previous diagonals are already filled in.
        for i in range(n - d + 1):
            j = i + d; # on the dth diagonal, j - i = d
            opt = i + 1 # until we find a better one
            for k in range(i+2, j+1):
                if C[i][k-1]+C[k][j] < C[i][opt-1]+C[opt][j]:
                    opt = k
            R[i][j] = opt
            W[i][j] = W[i][j-1] + a[j] + b[j]
            C[i][j] = C[i][opt-1] + C[opt][j] + W[i][j]
            # Draw this cell of the table in the given window.
            drawSquare(i, j, W[i][j], C[i][j], R[i][j], win, indent, squareSize )

def drawSquare(i,j, w, c, r, win, indent, squareSize):
    upperCorner = Point(j*squareSize+indent, (i)*squareSize+indent)
    upperCenter = Point(upperCorner.getX()+squareSize/2,upperCorner.getY()+indent)
    lowerCorner = Point(upperCorner.getX()+squareSize, upperCorner.getY()+squareSize )
    Rectangle(upperCorner, lowerCorner).draw(win)
    textR = Text(Point(upperCenter.getX(),upperCenter.getY()+5),"R"+str(i)+str(j)+": " + "%3d"%r)
    textW = Text(Point(upperCenter.getX(),upperCenter.getY()+30),"W"+str(i)+str(j)+": " + "%3d"%w)
    textC = Text(Point(upperCenter.getX(),upperCenter.getY()+55),"C"+str(i)+str(j)+": " + "%3d"%c)
    for t in [textR, textW, textC]:
        t.setFace("courier")
        t.draw(win)


# make an n-by-n list of all zeros
def makeTable(n):
    result = []
    for i in range(n+1):
        row = [0 for j in range(n+1)] 
        result += [row] 
    return result 

# display info about the optimal tree, based on the R list.
def optimalTree(R, letters):
    n = len(R) - 1
    print ("preorder:")
    preOrder(R, letters, 0, n)
    print ()
    print (" inorder:")
    inOrder(R, letters, 0, n)
    print ()

# display a preorder traversal of the tree    
def preOrder(R, letters, i, j):
    if i<j:
        k = R[i][j]
        print (letters[k-1], end=" ")
        preOrder(R, letters, i, k-1)
        preOrder(R, letters, k, j)

# display a preorder traversal of the tree       
def inOrder(R, letters, i, j):
    if i<j:
        k = R[i][j]
        inOrder(R, letters, i, k-1)
        print (letters[k-1], end=" ")
        inOrder(R, letters, k, j)

def main():          
    n = 5
    letters = ['A', 'E', 'I', 'O', 'U']
    a = [None, 32, 42, 26, 32, 12]
    b = [0, 34, 38, 58, 95, 21]
    
# An alternate setup for a different problem:
#    n = 3
#    letters = ['I', 'O', 'U']
#    p = [None, 6, 3, 6]
#    q = [7, 5, 2, 9]
    
    R = makeTable(n)   # initial table of all zeros.
    indent = 20
    squareSize = 100
    dimension = squareSize*(n+1)+2*indent
    win = GraphWin("Optimal tree", dimension, dimension)
    win.setBackground('white')
    optimalBST(n, a, b, R, win, indent, squareSize) 
    optimalTree(R, letters)

    # Keep the table display visible until the user clicks in the window.
    win.getMouse()
    win.close()

    
main()