zerosleeps

from pathlib import Path
import re
import unittest

def get_raw_input():
    return (Path(__file__).parent/'day_18_input.txt').read_text()

def parse_raw_input(raw_input):
    return [line.strip().replace(' ','') for line in raw_input.strip().splitlines()]

def arithmetic(expression, part=1):
    if type(expression) == re.Match:
        expression = expression[0]

    # For part two, when given an add expression and nothing else, simply
    # return the result of that add
    if part == 2 and re.match(r'^\d+\+\d+$', expression):
        return str(eval(expression))

    result = None
    operator = None

    # For part two, first deal with all the additions
    while part == 2 and re.search(r'\+', expression):
        expression = re.sub(r'\d+\+\d+', arithmetic, expression, 1)

    # Using re.findall to group concurrent digits into one "char"
    for char in re.findall(r'(\d+|[()+*])', expression):
        if char.isdecimal():
            if result is None:
                result = int(char)
            else:
                if operator == '+':
                    result += int(char)
                elif operator == '*':
                    result *= int(char)
                # Reset operator for next component of the expression
                operator = None
        elif char in ['+', '*']:
            operator = char

    return str(result)

def parse_expression(expression, part=1):
    regex = re.compile(r'\([\d\+\*]+\)')

    while re.search(regex, expression):
        expression = re.sub(regex, lambda e: arithmetic(e, part), expression, 1)

    return int(arithmetic(expression, part))

def part_one(raw_input):
    return sum([parse_expression(expression) for expression in parse_raw_input(raw_input)])

def part_two(raw_input):
    return sum([parse_expression(expression, part=2) for expression in parse_raw_input(raw_input)])

class TestPartOne(unittest.TestCase):
    def test_example_1(self):
        self.assertEqual(part_one("1 + 2 * 3 + 4 * 5 + 6"), 71)
    def test_example_2(self):
        self.assertEqual(part_one("1 + (2 * 3) + (4 * (5 + 6))"), 51)
    def test_example_3(self):
        self.assertEqual(part_one("2 * 3 + (4 * 5)"), 26)
    def test_example_4(self):
        self.assertEqual(part_one("5 + (8 * 3 + 9 + 3 * 4 * 3)"), 437)
    def test_example_5(self):
        self.assertEqual(part_one("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))"), 12240)
    def test_example_6(self):
        self.assertEqual(part_one("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"), 13632)

class TestPartTwo(unittest.TestCase):
    def test_example_1(self):
        self.assertEqual(part_two("1 + 2 * 3 + 4 * 5 + 6"), 231)
    def test_example_2(self):
        self.assertEqual(part_two("1 + (2 * 3) + (4 * (5 + 6))"), 51)
    def test_example_3(self):
        self.assertEqual(part_two("2 * 3 + (4 * 5)"), 46)
    def test_example_4(self):
        self.assertEqual(part_two("5 + (8 * 3 + 9 + 3 * 4 * 3)"), 1445)
    def test_example_5(self):
        self.assertEqual(part_two("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))"), 669060)
    def test_example_6(self):
        self.assertEqual(part_two("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"), 23340)

if __name__ == '__main__':
    print(f'Part one: {part_one(get_raw_input())}')
    print(f'Part two: {part_two(get_raw_input())}')