import json
import re
from itertools import combinations
from math import ceil
from tools.aoc import AOCDay
from typing import Any, List, Union

re_pair = re.compile(r"\[(\d+),(\d+)\]")
re_lnum = re.compile(r".*[^\d](\d+)")
re_rnum = re.compile(r"(\d+)")
re_bignum = re.compile(r"(\d\d+)")


class Snailfish:
    def __init__(self, fish_str: str):
        self.pairs = fish_str
        while self.reduce():
            pass

    def __add__(self, other: 'Snailfish'):
        return Snailfish("[%s,%s]" % (self.pairs, other.pairs))

    def reduce(self):
        open_count = 0
        for i, c in enumerate(self.pairs):
            if c == "]":
                open_count -= 1
            elif c == "[":
                open_count += 1
                if open_count == 5:
                    explode_pair_match = re_pair.search(self.pairs, i)
                    index_l, index_r = explode_pair_match.span()
                    explode_left, explode_right = map(int, explode_pair_match.groups())
                    lnum_old = rnum_old = lnum_new = rnum_new = ""
                    lnum_index = index_l
                    rnum_index = index_r

                    if left_num := re_lnum.search(self.pairs[:index_l]):
                        lnum_old = left_num.groups()[0]
                        lnum_new = str(int(lnum_old) + explode_left)
                        lnum_index = left_num.span()[1] - len(lnum_old)

                    if right_num := re_rnum.search(self.pairs, index_r):
                        rnum_old = right_num.groups()[0]
                        rnum_new = str(int(rnum_old) + explode_right)
                        rnum_index = right_num.span()[0]

                    self.pairs = "".join([
                        self.pairs[:lnum_index],
                        lnum_new,
                        self.pairs[lnum_index + len(lnum_old):index_l],
                        "0",
                        self.pairs[index_r:rnum_index],
                        rnum_new,
                        self.pairs[rnum_index + len(rnum_old):]
                    ])
                    return True

        if big_num_match := re_bignum.search(self.pairs):
            num = int(big_num_match.group())
            self.pairs = self.pairs.replace(big_num_match.group(), "[%d,%d]" % (num // 2, int(ceil(num / 2))), 1)
            return True

        return False

    def getMagnitude(self) -> int:
        pairs = self.pairs
        while match := re_pair.findall(pairs):
            for num1, num2 in match:
                pairs = pairs.replace("[%s,%s]" % (num1, num2), str(int(num1) * 3 + int(num2) * 2))

        return int(pairs)


class BinarySnailfish:
    def __init__(self, value: Union[int, List] = None, depth: int = 0, parent: 'BinarySnailfish' = None):
        self.fours = []
        self.depth = depth
        self.parent = parent
        if not isinstance(value, list):
            self.value = value
            self.left = None
            self.right = None
        else:
            self.value = None
            if isinstance(value[0], BinarySnailfish):
                self.left = value[0]
                self.left.setDepth(self.depth + 1)
                self.left.parent = self
            else:
                self.left = BinarySnailfish(value[0], parent=self, depth=self.depth + 1)

            if isinstance(value[1], BinarySnailfish):
                self.right = value[1]
                self.right.setDepth(self.depth + 1)
                self.right.parent = self
            else:
                self.right = BinarySnailfish(value[1], parent=self, depth=self.depth + 1)

    def setDepth(self, depth: int):
        self.depth = depth
        if self.value is not None:
            return
        else:
            if self.depth == 4:
                self.getRoot().fours.append(self)
            self.left.setDepth(depth + 1)
            self.right.setDepth(depth + 1)

    def getRoot(self) -> 'BinarySnailfish':
        root = self
        while root.parent:
            root = root.parent

        return root

    def getLeftNumber(self) -> Union[None, 'BinarySnailfish']:
        parent = self.parent
        if self == self.parent.left:
            while parent.parent and parent == parent.parent.left:
                parent = parent.parent
            parent = parent.left
        else:
            while parent.parent and parent == parent.parent.right:
                parent = parent.parent
            parent = parent.right
        while parent.right.value is None:
            parent = parent.right
        return parent

        leftnum = None
        for current in self.getRoot().traverse_preorder():
            if current == self:
                return leftnum
            if current.value is not None:
                leftnum = current

    def getRightNumber(self) -> Union[None, 'BinarySnailfish']:
        returnNext = 0
        for current in self.getRoot().traverse_preorder():
            if returnNext >= 3 and current.value is not None:
                return current
            if returnNext > 0:
                returnNext += 1
            if current == self:
                returnNext += 1

    def split(self):
        self.left = BinarySnailfish(self.value // 2, depth=self.depth + 1, parent=self)
        self.right = BinarySnailfish(int(ceil(self.value / 2)), depth=self.depth + 1, parent=self)
        self.value = None

    def explode(self):
        left_num = self.getLeftNumber()
        if left_num is not None and left_num.value is not None:
            left_num.value += self.left.value

        right_num = self.getRightNumber()
        if right_num is not None and right_num.value is not None:
            right_num.value += self.right.value

        self.left = self.right = None
        self.value = 0

    def reduce(self):
        if self.value is not None:
            return
        found = True
        while found:
            found = False
            for n in self.traverse_preorder():
                if n.value is None and n.depth == 4:
                    n.explode()
                    found = True
                    break

            if found:
                continue

            for n in self.traverse_preorder():
                if n.value is not None and n.value > 9:
                    n.split()
                    found = True
                    break

    def traverse_preorder(self):
        yield self

        if self.left:
            for x in self.left.traverse_preorder():
                yield x

        if self.right:
            for x in self.right.traverse_preorder():
                yield x

    def __add__(self, other):
        sum_fish = BinarySnailfish([self, other])
        sum_fish.reduce()
        return sum_fish

    def __str__(self):
        if self.value is None:
            return "[%s,%s]" % (str(self.left), str(self.right))
        else:
            return str(self.value)

    def getMagnitude(self):
        if self.value is not None:
            return self.value
        else:
            return self.left.getMagnitude() * 3 + self.right.getMagnitude() * 2


class Day(AOCDay):
    test_solutions_p1 = [4140, 4417]
    test_solutions_p2 = [3993, 4796]

    def part1(self) -> Any:
        #snailfishes = [Snailfish(x) for x in self.getInput()]
        snailfishes = [BinarySnailfish(json.loads(x)) for x in self.getInput()]
        return sum(snailfishes[1:], start=snailfishes[0]).getMagnitude()

    def part2(self) -> Any:
        snailfishes = [Snailfish(x) for x in self.getInput()]
        max_mag = 0
        for a, b in combinations(snailfishes, 2):
            sub_mag_a = (a + b).getMagnitude()
            sub_mag_b = (b + a).getMagnitude()
            max_mag = max(max_mag, sub_mag_a, sub_mag_b)

        return max_mag