aoc_2021_18

def parse(s):
    '''
    Parse string into a list of list...
    string e.g. '[1,2]'
    '''
    if s.isnumeric():
        return int(s)

    ptr = 1
    open_brack = 0
    v = []
    while ptr < len(s)-1:
        if s[ptr].isnumeric() and open_brack == 0:
            v.append(s[ptr])
        elif s[ptr] == '[':
            open_brack += 1
            if open_brack == 1:
                begin = ptr
        elif s[ptr] == ']':
            open_brack -= 1
            if open_brack == 0:
                end = ptr
                v.append(s[begin:end+1])

        ptr += 1

    return [parse(v[0]), parse(v[1])]

def explode(arr, depth):
    '''
    if the bracket is surrounded by 4 outer bracket
    it will be reduced to 0, and distribute the left and
    right values out to the nearest int
    '''
    if type(arr) == list and depth==4:
        left_v = arr[0]
        right_v = arr[1]
        return 0, left_v, right_v

    if type(arr) == int:
        return arr, 0, 0

    update = [0, 0]
    for i in range(2):
        if type(arr[i]) == list:
            arr[i], lv, rv= explode(arr[i], depth+1)
            vs = [lv, rv]
            # the left value of a left node will be sent upward
            # same for the right value of a right node
            update[i] = vs[i]

            # If there exist a neighbour, distribute the other
            # value to the neighbour
            if vs[(i+1)%2] > 0:
                if type(arr[(i+1)%2]) == int:
                    arr[(i+1)%2] += vs[(i+1)%2]
                else:
                    b = arr[(i+1)%2]
                    while type(b) != int:
                        if type(b[i]) == int:
                            b[i] += vs[(i+1)%2]
                            break
                        else:
                            b = b[i]

    return arr, update[0], update[1]

def split(arr):
    '''
    Only perform on the left-most >=10 integer
    '''
    if type(arr)==int:
        if arr>=10:
            return [arr//2, arr//2+arr%2], True
        else:
            return arr, False

    for i in range(2):
        arr[i], splited = split(arr[i])
        if splited:
            return arr, splited
    return arr, splited

def depth(arr):
    if type(arr) == int:
        return 0
    return max(depth(arr[0])+1, depth(arr[1])+1)

def value(arr):
    if type(arr) == int:
        if arr >= 10:
            return False
        else:
            return True
    return value(arr[0]) and value(arr[1])

def snail_sum(arr):
    if type(arr) == int:
        return arr

    return 3*snail_sum(arr[0])+2*snail_sum(arr[1])

####################################################
with open('input_day18.txt') as f:
    snail = f.readlines()
    for i, v in enumerate(snail):
        snail[i] = snail[i].strip()

# part 1, find the sum....
s0 = parse(snail[0])
for s in snail[1:]:
    s1 = parse(s)
    new_s = [s0, s1]
    #print(new_s)
    while depth(new_s)>=5 or not value(new_s):
        new_s = explode(new_s, 0)[0]
        new_s = split(new_s)[0]

    s0 = new_s
print(snail_sum(s0))

# part 2, largest sum?
max_sum = 0
for i in range(len(snail)):
    for j in range(len(snail)):
        if i == j:
            continue
        s0 = parse(snail[i])
        s1 = parse(snail[j])
        new_s = [s0, s1]
        while depth(new_s)>=5 or not value(new_s):
            new_s = explode(new_s, 0)[0]
            new_s = split(new_s)[0]
        max_sum = max(max_sum, snail_sum(new_s))
print(max_sum)