tree.py

'''
AICogSci Final Project

base tree code from http://knuth.luther.edu/~leekent/CS2Plus/chap6/chap6.html
Added:
- EquationTree class to generate a specialized tree
- Branch and Equals Node classes to make printing more legible
- canEval, getInverse and inorder methods to most of the nodes to allow further functionality
- VariableNode class was added to account for nodes with variables
- add_to_tree and generate_tree methods used to create and change equation tree class'''

import queue

class EquationTree:
    def __init__(self):
        self.root = EqualsNode()
        self.root.left = None
        self.root.right = None

    def updateNode(self, parent, left,  newNode):
        if left: parent.left = newNode
        else: parent.right = newNode
        self.printTree()

    def get_expression(self):
        self.output = ""
        self.inorder(self.root)
        print("output is: ", self.output)
        return self.output

    def inorder(self, node):
        if node:
            self.inorder(node.left)
            self.output += str(node.value)
            self.inorder(node.right)

    def solve(self):
        # solves equation assuming there are only a variable node and num node left
        if isinstance(self.root.left, VariableNode):
            # variablenode is on left of equation
            sol = self.root.right.eval() / self.root.left.eval()
            self.root.left = VariableNode('x')
            self.root.right = NumNode(int(sol))
        else:
            sol = self.root.left.eval() / self.root.right.eval()
            self.root.right = VariableNode('x')
            self.root.left = NumNode(int(sol))
        return self.get_expression()

    def evaluate(self, parent, left):
        # can simplify this to call helper function
        if left:
            if parent.left.canEval():
                print("parent.left.left is ", parent.left.left.value)
                print("parent.left.right is ", parent.left.right)
                if isinstance(parent.left.left, NumNode):
                    print("num node")
                    print("evaluating", parent.left.eval())
                    newNode = NumNode(parent.left.eval())
                else:
                    print("variable node")
                    print("evaluating", parent.left.eval())
                    newNode = VariableNode(str(parent.left.eval()) + "x")
                parent.left = newNode
            elif parent.left.left != None and parent.left.left.canEval():
                self.evaluate(parent.left, True)
            elif parent.left.right != None and parent.left.right.canEval():
                self.evaluate(parent.left, False)
        else:
            if parent.right.canEval():
                if isinstance(parent.right.left, NumNode):
                    print("evaluating", parent.right.eval())
                    newNode = NumNode(parent.right.eval())
                else:
                    print("evaluating", parent.right.eval())
                    newNode = VariableNode(str(parent.right.eval()) + "x")
                parent.right = newNode
            elif parent.right.left != None and parent.right.left.canEval():
                self.evaluate(parent.right, True)
            elif parent.right.right != None and parent.right.right.canEval():
                self.evaluate(parent.right, False)

    def printTree(self):
        q = queue.Queue()
        q.enqueue(self.root)
        while not q.isEmpty():
            currentNode = q.front()
            if currentNode.left != None:
                b = Branch()
                q.enqueue(b)
                q.enqueue(currentNode.left)
            if currentNode.right != None:
                q.enqueue(currentNode.right)
            print(q.dequeue().value, end ='     ')

class Branch:
    def __init__(self):
        self.value = "\n"
        self.left = None
        self.right = None

class EqualsNode:
    def __init__(self):
        self.value = "="


class TimesNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
        self.value = "*"
        
    
    def eval(self):  
        return self.left.eval() * self.right.eval()
    
    def inorder(self):
        return "(" + self.left.inorder() + " * " + self.right.inorder() + ")" 
    
class PlusNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
        self.value = "+"
        
    def canEval(self):
        if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
            return True
        elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode):
            return True
        else:
            return False

    def getInverse(self):
        return MinusNode(None,None)

    def eval(self):
        return self.left.eval() + self.right.eval()
    
    
    def inorder(self):
        return "(" + self.left.inorder() + " + " + self.right.inorder() + ")"  

class MinusNode:
    def __init__(self,left, right):
        self.left = left
        self.right = right
        self.value = "-"

    def canEval(self):
        if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
            return True
        elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode):
            return True
        else:
            return False

    def getInverse(self):
        return PlusNode(None, None)

    def eval(self):
        return self.left.eval() - self.right.eval()

    def inorder(self):
        return "(" + self.left.inorder() + "-" + self.right.inorder() + ")"

class NumNode:
    def __init__(self, num):
        self.value = num
        self.left = None
        self.right = None
        
    def eval(self):
        return self.value

    def canEval(self):
        return False # not an operator
    
    def inorder(self):
        return str(self.value)

class VariableNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def eval(self):
        # get value before x
        return int(self.value.split("x")[0])

    def canEval(self):
        return False 

    def inorder(self):
        return str(self.eval())

OPERATORS = set(['+', '-', '*', '/', '(', ')'])
PRIORITY = {'+':1, '-':1, '*':2, '/':2}

def infix_to_prefix(formula):
    # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py
    op_stack = []
    exp_stack = []
    for ch in formula.split():
        if not ch in OPERATORS:
            exp_stack.append(ch)
        elif ch == '(':
            op_stack.append(ch)
        elif ch == ')':
            while op_stack[-1] != '(':
                op = op_stack.pop()
                a = exp_stack.pop()
                b = exp_stack.pop()
                exp_stack.append( op+b+a )
            op_stack.pop() # pop '('
        else:
            while op_stack and op_stack[-1] != '(' and PRIORITY[ch] <= PRIORITY[op_stack[-1]]:
                op = op_stack.pop()
                a = exp_stack.pop()
                b = exp_stack.pop()
                exp_stack.append( op+" "+b+" "+a )
            op_stack.append(ch)
    while op_stack:
        op = op_stack.pop()
        a = exp_stack.pop()
        b = exp_stack.pop()
        exp_stack.append( op+" "+b+" "+a )
    print(exp_stack[-1])
    return exp_stack[-1]

def add_to_tree(q, root, left):
    if not q.isEmpty():
        token = q.dequeue()
        operator = False
        if token == "+":
            newNode = PlusNode(None, None)
            operator = True
        elif token == "*":
            newNode = TimesNode(None, None)
            operator = True

        elif token == "-":
            newNode = MinusNode(None, None)
            operator = True
        
        elif "x" in token:
            newNode = VariableNode(token)
        else:
            newNode = NumNode(int(token)) 

        if left: root.left = newNode
        else:  root.right = newNode
        if operator:
            add_to_tree(q, newNode, True)
            add_to_tree(q, newNode, False)
        return newNode
        

def generate_tree(equation):
    t = EquationTree()

    expressions = equation.split("=")

    x1 = infix_to_prefix(expressions[0])
    x2 = infix_to_prefix(expressions[1])

    lst1 = x1.split()
    q1 = queue.Queue()

    lst2 = x2.split()
    q2 = queue.Queue()    

    for token in lst1:
        q1.enqueue(token)
    add_to_tree(q1, t.root, True) # add expression tree to left side of equation tree

    for token in lst2:
        q2.enqueue(token)
    add_to_tree(q2, t.root, False) # add expression tree to right side of equation tree

    t.printTree()
    return t