import aoclib
import re

DAY = 18
TEST_SOLUTION_PART1 = 26457
TEST_SOLUTION_PART2 = 694173

r_parentheses = re.compile(r'(\([^()]+\))')
r_calc = re.compile(r'^(\d+([*+])\d+)')
r_add = re.compile(r'[^+0-9]?(\d+\+\d+)')
r_mul = re.compile(r'[^*0-9]?(\d+\*\d+)')


def calculate(expression, add_pred=False):
    while "(" in expression:
        for parentheses in re.findall(r_parentheses, expression):
            expression = expression.replace(parentheses, str(calculate(parentheses[1:-1], add_pred)), 1)

    if not add_pred:
        while "+" in expression or "*" in expression:
            calc = re.match(r_calc, expression)[0]
            expression = expression.replace(calc, str(eval(calc)), 1)
    else:
        while "+" in expression:
            for addition in re.findall(r_add, expression):
                left, right = addition.split('+')
                expression = re.sub(
                    "(^|[*+])%s($|[+*])" % addition.replace('+', '\+'),
                    "\g<1>%s\g<2>" % (int(left) + int(right)),
                    expression
                )

        while "*" in expression:
            for multiplication in re.findall(r_mul, expression):
                left, right = multiplication.split('*')
                expression = re.sub(
                    "(^|[*+])%s($|[+*])" % multiplication.replace('*', '\*'),
                    "\g<1>%s\g<2>" % (int(left) * int(right)),
                    expression
                )

    return int(expression)


def part1(test_mode=False, add_pred=False):
    my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)

    result = 0
    for line in my_input:
        line = line.replace(" ", "")
        result += calculate(line, add_pred)

    return result


def part2(test_mode=False):
    return part1(test_mode, add_pred=True)