AoC 21-18

from abc import ABC, abstractmethod
import itertools

class Number(ABC):
    def __init__(self) -> None:
        self.parent = None

    def _check_explosions(self, nesting_level: int):
        """Should handle any explosions and return the current Node (or its replacement) and a boolean indicating if
        the node (or any of its children) exploded

        Parameters
        ----------
        nesting_level : int
            How deep the current node is in the tree

        Returns
        -------
        Number, bool
            The current node (or its replacement) and whether the node or its children exploded.
        """
        return self, False

    @abstractmethod
    def _check_splits(self):
        """Should handle any splits and return the current Node (or its replacement) along with a boolean indicating if
        the node or any of its children did split.
        """
        pass

    @abstractmethod
    def magnitude(self):
        pass

    def reduce(self):
        pass

class SnailfishNumber(Number):
    def __init__(self, raw: str) -> None:
        super().__init__()
        self.left = None
        self.right = None

        # Parse raw number
        self.raw = raw
        if self.raw is not None:
            self._parse_raw_number()

    def __repr__(self) -> str:
        return f"[{self.left}, {self.right}]"

    def _split_pair(raw: str):
        """Splits a string SnailfishNumber into the left and right children. Used for parsing the raw snailfish strings.

        Parameters
        ----------
        raw : str
            The raw snailfish string

        Returns
        -------
        str, str
            The left and right children in string form
        """
        bracket_count = 0
        for idx, char in enumerate(raw):
            if char == "[":
                bracket_count += 1
            elif char == "]":
                bracket_count -= 1

            if (bracket_count == 1) and (idx != 0):
                # Found matching bracket to inner pair
                return raw[1:idx + 1], raw[idx + 2: -1]


    def _parse_raw_number(self):
        """Parses a raw string as a SnailfishNumber
        """
        # Get rid of brackets
        left, right = SnailfishNumber._split_pair(self.raw)

        self.left = SnailfishNumber(left) if left[0] == "[" else HumanNumber(int(left))
        self.right = SnailfishNumber(right) if right[0] == "[" else HumanNumber(int(right))

        self.left.parent = self
        self.right.parent = self


    def add(self, other):
        """Adds another SnailfishNumber to the current SnailfishNumber.

        Parameters
        ----------
        other : SnailfishNumber
            The number to add.

        Returns
        -------
        SnailfishNumber
            The reduced result of the snailfish addition.
        """
        sn = SnailfishNumber(None)
        sn.left = self
        sn.right = other

        self.parent = sn
        other.parent = sn

        sn.reduce()

        return sn

    def reduce(self):
        """Reduces the current node if any reductions need to be performed. This is an in-place operation.
        """
        did_explode = True
        did_split = True

        while did_explode or did_split:
            did_explode = False
            did_split = False

            _, did_explode = self._check_explosions(0)

            if did_explode:
                continue

            _, did_split = self._check_splits()

    def magnitude(self):
        """Calculates the magnitude of the current node

        Returns
        -------
        int
            The node's magnitude
        """
        left_magnitude = self.left.magnitude()
        right_magnitude = self.right.magnitude()

        return 3 * left_magnitude + 2 * right_magnitude

    def _check_splits(self):
        """Determines and handles any splits that occur due to large HumanNumbers (>= 10)

        Returns
        -------
        Number, bool
            The current node along with whether any of its children split
        """
        self.left, left_did_split = self.left._check_splits()
        self.left.parent = self

        if left_did_split:
            return self, True

        self.right, right_did_split =  self.right._check_splits()
        self.right.parent = self

        return self, right_did_split

    def _find_deepest_directional_child(self, direction: str):
        """Finds the deepest child in a given consistent direction from the current node.

        Parameters
        ----------
        direction : str
            Which direction to search in.

        Returns
        -------
        HumanNumber
            The deepest HumanNumber node to the direction of the current node.
        """
        assert direction in ("left", "right"), "Direction must either be 'left' or 'right'."

        current = self
        while type(current) is SnailfishNumber:
            current = getattr(current, direction)

        return current

    def _find_nearest_directional_leaf(self, direction: str):
        """Finds the nearest leaf node either to the left or the right of the current node anywhere in the tree.

        Parameters
        ----------
        direction : str
            Which direction to search in (either left or right)

        Returns
        -------
        HumanNumber
            The HumanNumber node closest to the current node in the given direction
        """
        assert direction in ("left", "right"), "Direction must either be 'left' or 'right'."
        anti_direction = "right" if direction == "left" else "left"
        current = self
        first_directional_child = None

        while current.parent is not None and first_directional_child is None:
            directional_child = getattr(current.parent, direction)
            if directional_child is current:
                current = current.parent
            else:
                if type(directional_child) is HumanNumber:
                    # Found it
                    return directional_child
                else:
                    first_directional_child = directional_child._find_deepest_directional_child(anti_direction)

        return first_directional_child

    def _check_explosions(self, nesting_level: int):
        """Checks if the current node will explode and handles the explosion if so.

        Parameters
        ----------
        nesting_level : int
            The current node's nesting level within the tree.

        Returns
        -------
        Number, bool
            The current node (or its replacement) along with whether the node or its children exploded
        """
        if nesting_level < 4:
            # Not nested deep enough so we check the children
            new_left, left_did_explode = self.left._check_explosions(nesting_level + 1)
            new_left.parent = self
            self.left = new_left

            if left_did_explode:
                return self, True

            new_right, right_did_explode = self.right._check_explosions(nesting_level + 1)
            new_right.parent = self
            self.right = new_right

            return self, right_did_explode

        else:
            # Explosion
            # Find first left-child
            first_left_child = self._find_nearest_directional_leaf("left")
            if first_left_child is not None:
                first_left_child.value += self.left.value

            # Find first right-child
            first_right_child = self._find_nearest_directional_leaf("right")
            if first_right_child is not None:
                first_right_child.value += self.right.value

            return HumanNumber(0), True

class HumanNumber(Number):
    def __init__(self, value: int) -> None:
        super().__init__()
        self.value = int(value)

    def __repr__(self) -> str:
        return f"{self.value}"

    def _check_splits(self):
        """Handles a split of a HumanNumber

        Returns
        -------
        Number, bool
            Either the current node or its split replacement, along with whether the node split
        """
        if self.value >= 10:
            sn = SnailfishNumber(None)
            sn.left = HumanNumber(self.value // 2)
            sn.right = HumanNumber(self.value - self.value // 2)
            sn.parent = self.parent
            sn.left.parent = self
            sn.right.parent = self

            return sn, True

        return self, False

    def magnitude(self):
        return self.value

if __name__ == "__main__":
    file = "inputs/2021/18/data.txt"

    with open(file, "r") as f:
        raw_numbers = f.read().split("\n")

    snailfish_sum = SnailfishNumber(raw_numbers[0])
    for raw_number in raw_numbers[1:]:
        second_number = SnailfishNumber(raw_number)
        snailfish_sum = snailfish_sum.add(second_number)

    print(f"Part 1: {snailfish_sum.magnitude():,d}")

    max_magnitude = -1
    for pair in itertools.combinations(raw_numbers, 2):
        a = SnailfishNumber(pair[0])
        b = SnailfishNumber(pair[1])
        max_magnitude = max(max_magnitude, a.add(b).magnitude())

    print(f"Part 2: {max_magnitude:,d}")