# Algorithm to calculate the best stocks to buy based on their forecasts and taking into account transaction costs when
# selling one stock for another.

stock_predictions = {
    'APPL': [1, 1, 1, 10, 1, 1],
    'GOOG': [0, 0, 0, 0, 0, 10],
    'MSFT': [0, 0, 10, 0, 0, 0]
}


class Node:
    def __init__(self, stock, day, unit_value, total_value, parent):
        self.stock = stock
        self.day = day
        self.unit_value = unit_value
        self.total_value = total_value
        self.parent = parent
        self.children = []

    def __str__(self):
        return 'Stock: ' + self.stock + ' on day ' + str(self.day) + ' trading at ' + str(
            self.unit_value) + '. Total Value: ' + str(self.total_value)


def add_nodes(node, stocks, day, transaction_cost):
    length = len(next(iter(stocks.values())))

    if day < length:
        for stock_name, predictions in stock_predictions.items():
            cost = transaction_cost if stock_name != node.stock else 0
            total = node.total_value + predictions[day] - cost
            child = Node(stock_name, day, predictions[day], total, node)
            node.children.append(child)
            add_nodes(node.children[-1], stocks, day + 1, transaction_cost)


def print_nodes(parent):
    print(str(parent.value))

    if len(parent.children) > 0:
        for child in parent.children:
            print_nodes(child)


def get_greatest_node(root):
    max_node = root
    if len(max_node.children) > 1:
        for child in root.children:
            max_value = max_node.total_value
            greatest_child = get_greatest_node(child)
            greatest_child_value = greatest_child.total_value
            if greatest_child_value > max_value:
                max_node = greatest_child

    return max_node


def print_path(node):
    if node.stock == 'root':
        return

    print(str(node))
    return print_path(node.parent)


if __name__ == "__main__":
    transaction_cost = 5
    root = Node('root', 0, 0, transaction_cost, None)
    add_nodes(root, stock_predictions, 0, transaction_cost)
    greatest_node = get_greatest_node(root)
    print_path(greatest_node)