import shared
import itertools
import json
import math
from functools import reduce

red = lambda word: f"\u001b[31m{word}\u001b[0m"
yellow = lambda word: f"\u001b[33m{word}\u001b[0m"
blue = lambda word: f"\u001b[34m{word}\u001b[0m"
green = lambda word: f"\u001b[32m{word}\u001b[0m"

class Snailfish:
    def __init__(self):
        self.load()

    def load(self):
        snailfish = []
        with open(shared.get_fname(18), "r") as f:
            for line in f.readlines():
                snailfish.append(json.loads(line.rstrip()))
                print(snailfish[-1])
        self.snailfish = snailfish

    def run(self):
        reduced = reduce(self.add, self.snailfish)
        print(green(reduced))
        self.magnitude_sum = self.magnitude(reduced)
        permutations = itertools.permutations(self.snailfish, 2)
        self.magnitude_max = 0
        for a,b in permutations:
            print(red(a), yellow(b))
            fish_sum = self.add(a,b)
            print("\t", blue(fish_sum))
            mag = self.magnitude(fish_sum)
            if mag > self.magnitude_max:
                self.magnitude_max = mag
            print()
        
    def add_left(self, fish, value):
        if value is None:
            return fish
        if isinstance(fish, int):
            return fish + value
        return [self.add_left(fish[0], value), fish[1]]

    def add_right(self, fish, value):
        if value is None:
            return fish
        if isinstance(fish, int):
            return fish + value
        return [fish[0], self.add_right(fish[1], value)]

    def explode(self, fish, depth=4):
        if isinstance(fish, int):
            return False, None, fish, None
        if depth == 0:
            return True, fish[0], 0, fish[1]
        x, y = fish
        exploded, left, x, right = self.explode(x, depth - 1)
        if exploded:
            return True, left, [x, self.add_left(y, right)], None
        exploded, left, y, right = self.explode(y, depth - 1)
        if exploded:
            return True, None, [self.add_right(x, left), y], right
        return False, None, fish, None

    def split(self, fish):
        if isinstance(fish, int):
            if fish >= 10:
                return True, [fish // 2, math.ceil(fish / 2)]
            return False, fish
        x, y = fish
        split, x = self.split(x)
        if split:
            return True, [x, y]
        split, y = self.split(y)
        return split, [x, y]

    def add(self, x, y):
        fish = [x, y]
        while True:
            exploded, _, fish, _ = self.explode(fish)
            if exploded:
                continue
            split, fish = self.split(fish)
            if not split:
                break
        return fish

    def magnitude(self, fish):
        if isinstance(fish, int):
            return fish
        return 3 * self.magnitude(fish[0]) + 2 * self.magnitude(fish[1])


def main():
    s = Snailfish()
    s.run()
    print('part1: sum', s.magnitude_sum)
    print('part2: max', s.magnitude_max)


if __name__ == "__main__":
    main()