from itertools import combinations
from math import ceil
from tools.aoc import AOCDay
from typing import Any


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

    def __add__(self, other):
        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:
                    index_explode = i
                    while self.pairs[index_explode + 1] == "[" or self.pairs[index_explode + 3] == "[":
                        if self.pairs[index_explode + 1] == "[":
                            index_explode += 1
                        if self.pairs[index_explode + 3] == "[":
                            index_explode += 3
                    explode_pair = self.pairs[index_explode:self.pairs.index("]", index_explode) + 1]
                    explode_left, explode_right = map(int, explode_pair[1:-1].split(","))
                    old_number_left = ""
                    old_number_right = ""
                    new_number_left = ""
                    new_number_right = ""
                    index_left = index_explode
                    index_right = index_explode + len(explode_pair)
                    for i_left in range(index_explode, 0, -1):
                        if self.pairs[i_left] not in ['[', ']', ',']:
                            il_add = 0
                            while self.pairs[i_left - il_add - 1] in "0123456789":
                                il_add += 1
                            index_left = i_left - il_add
                            old_number_left = self.pairs[i_left - il_add:i_left + 1]
                            new_number_left = str(int(old_number_left) + explode_left)
                            break
                    for i_right in range(index_explode + len(explode_pair), len(self.pairs)):
                        if self.pairs[i_right] not in ['[', ']', ',']:
                            ir_add = 0
                            while self.pairs[i_right + ir_add + 1] in "0123456789":
                                ir_add += 1
                            index_right = i_right
                            old_number_right = self.pairs[i_right:i_right + ir_add + 1]
                            new_number_right = str(int(old_number_right) + explode_right)
                            break
                    new_pairs = ""
                    if index_left != index_explode:
                        new_pairs += self.pairs[:index_left] + new_number_left + self.pairs[index_left + len(old_number_left):index_explode]
                    else:
                        new_pairs = self.pairs[:index_explode]
                    new_pairs += "0"
                    if index_right != index_explode + len(explode_pair):
                        new_pairs += self.pairs[index_explode + len(explode_pair):index_right] + new_number_right + self.pairs[index_right + len(old_number_right):]
                    else:
                        new_pairs += self.pairs[index_explode + len(explode_pair):]
                    self.pairs = new_pairs
                    return True

        for i in range(1, len(self.pairs)):
            if self.pairs[i - 1] not in ["[", "]", ","] and self.pairs[i] not in ["[", "]", ","]:
                number = int(self.pairs[i-1:i+1])
                self.pairs = self.pairs.replace(str(number), "[%d,%d]" % (number // 2, int(ceil(number / 2))), 1)
                return True

        return False

    def getMagnitude(self) -> int:
        pairs = self.pairs
        while "[" in pairs:
            for i in range(len(pairs)):
                if pairs[i] == "]":
                    r = i + 1
                    l = i - 1
                    while pairs[l] != "[":
                        l -= 1
                    replace_pair = pairs[l:r]
                    num1, num2 = map(int, replace_pair[1:-1].split(","))
                    pairs = pairs.replace(replace_pair, str(num1 * 3 + num2 * 2))
                    break

        return int(pairs)


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()]
        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