AOC - Python - Day 18

"""
--- Day 18: Operation Order ---
"""

import re
import doctest

from aocd import get_data  # https://github.com/wimglenn/advent-of-code-data
day = 18
year = 2020

def match_parens(str):
    """
    Get all the parenthesis depths for a given string
    """
    depth = 0
    tmp = [0,] * len(str)
    for index, char in enumerate(str):
        if char == "(":
            depth += 1
            tmp[index] = depth
        elif char == ")":
            tmp[index] = depth
            depth -= 1
    return tmp

def calculate(line, advanced=False):
    """
    Calculate the elf math total of the provided line
    :param line: the line as read from the inout file
    :param advanced: should addition happen before multiplication?
    :return: the calculated total
    """
    parens = match_parens(line)
    depth = max(parens)
    if depth == 0:
        return calculate_left_to_right(line, advanced)
    else:
        depth_index = parens.index(max(parens))
        closing = parens[depth_index+1:].index(depth)
        value = calculate_left_to_right(line[depth_index+1: depth_index+closing+1], advanced)
        str_replace = line[depth_index: depth_index+closing+2]
        line = line.replace(str_replace, str(value))
        return calculate(line, advanced)


def calculate_left_to_right(eq, advanced=False):
    """
    Calculate a given sequence as either strict left to right or if advanced than left to right with addition
    being performed first.
    :param eq: the immediate equation to be calculated
    :param advanced: addition performed BEFORE multiplication?
    :return: the calculated value
    """
    total = 0

    numbers = [int(x) for x in re.split(r"\*|\+", eq)]  # get all the numbers
    operands = [x for x in re.split(r"[0-9]", eq) if x != '']   # get all the operands

    if not advanced:
        total = numbers[0]
        for i in range(1, len(numbers)):
            operand = operands[i-1]
            right = numbers[i]
            if operand == "+":
                total += right
            elif operand == "*":
                total *= right
        return total

    else:
        # Calculate left to right EXCEPT with + first
        total = 0

        while "+" in operands:
            numbers, operands = find_perform_addition(numbers, operands)

        if len(operands) == 0:
            total = numbers[0]
            return total

        else:
            total = numbers[0]
            for i in range(1, len(numbers)):
                total *= numbers[i]
            return total


def find_perform_addition(numbers, operands):
    numbers_copy = numbers.copy()
    operands_copy = operands.copy()

    for index, op in enumerate(operands):
        if op == "+":
            x = numbers_copy.pop(index)
            y = numbers_copy.pop(index)
            numbers_copy.insert(index, x+y)
            operands_copy.pop(index)
            return (numbers_copy, operands_copy)
    return (numbers_copy, operands_copy)


def part_a(data):
    """
    >>> line = "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2".replace(" ", "")
    >>> calculate(line)
    13632
    """

    total = 0
    for line in data:
        total += calculate(line.replace(" ", ""))
    print("Part A total: {:d}".format(total))


def part_b(data):
    """
    >>> line = "1 + (2 * 3) + (4 * (5 + 6))".replace(" ", "")
    >>> calculate(line, True)
    51
    >>> line = "5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))".replace(" ", "")
    >>> calculate(line, True)
    669060
    >>> line = "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2".replace(" ", "")
    >>> calculate(line, True)
    23340
    """

    total = 0
    for line in data:
        total += calculate(line.replace(" ", ""), True)
    print("Part B total: {:d}".format(total))


if __name__ == "__main__":
    doctest.testmod()

    data = get_data(day=day, year=year).splitlines()
    part_a(data)
    part_b(data)