from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass
from functools import reduce
from itertools import combinations
from math import floor, ceil
from utils import open_day

@dataclass
class Pair:
    car: int | Pair
    cdr: int | Pair
    def __repr__(self) -> str:
        return f'[{self.car},{self.cdr}]'

def car(p, assign=None):
    if assign is not None:
        p.car = assign
    return p.car
def cdr(p, assign=None):
    if assign is not None:
        p.cdr = assign
    return p.cdr

def is_pair(p): return isinstance(p, Pair)

def parse(s): return eval(s.replace('[', 'Pair(').replace(']', ')'))

with open_day(18) as f:
    homework = [parse(l) for l in f]

def try_explode(n):
    def try_add(n, val, access, primary, secondary):
        if val == 0: return 0
        content = access(n)
        if not is_pair(content):
            access(n, content + val)
            return 0
        val = try_add(content, val, primary, primary, secondary)
        val = try_add(content, val, secondary, primary, secondary)
        return val
    def descend(n, depth=0):
        if not is_pair(n): return None
        if depth == 4: return n
        car_val = descend(n.car, depth + 1)
        if car_val is not None:
            if depth == 3: n.car = 0
            return Pair(car_val.car, try_add(n, car_val.cdr, cdr, car, cdr))
        else:
            cdr_val = descend(n.cdr, depth + 1)
            if cdr_val is not None:
                if depth == 3: n.cdr = 0
                return Pair(try_add(n, cdr_val.car, car, cdr, car), cdr_val.cdr)
    return descend(n)

tests = [
    ('[[[[[9,8],1],2],3],4]', '[[[[0,9],2],3],4]'),
    ('[7,[6,[5,[4,[3,2]]]]]', '[7,[6,[5,[7,0]]]]'),
    ('[[6,[5,[4,[3,2]]]],1]', '[[6,[5,[7,0]]],3]'),
    ('[[3,[2,[1,[7,3]]]],[6,[5,[4,[3,2]]]]]', '[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]'),
    ('[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]', '[[3,[2,[8,0]]],[9,[5,[7,0]]]]'),
]

for test, expect in tests:
    test = parse(test)
    try_explode(test)
    assert(test == parse(expect))

def try_split(n, access=lambda n: n):
    content = access(n)
    if is_pair(content):
        return try_split(content, car) or try_split(content, cdr)
    if content < 10: return False
    content /= 2
    access(n, Pair(floor(content), ceil(content)))
    return True

def reduce_num(n):
    n = deepcopy(n)
    while True:
        if try_explode(n): continue
        if try_split(n): continue
        break
    return n

def magnitude(n):
    if is_pair(n):
        return 3 * magnitude(n.car) + 2 * magnitude(n.cdr)
    return n

print(magnitude(reduce(lambda a, b: reduce_num(Pair(a, b)), homework)))
print(max(max(magnitude(reduce_num(Pair(a, b))), magnitude(reduce_num(Pair(b, a)))) for a, b in combinations(homework, 2)))