import re


def decode_line(line):
    mask = re.search(r"mask = (?P<mask>[1 | 0 | X]{36})", line)
    mem = re.search(r"mem\[(?P<address>\d+)] = (?P<value>\d+)", line)
    if mask:
        return 'm', mask.group('mask')
    if mem:
        return mem.group('address', 'value')


def to_int(num):
    new = 0
    for j in range(0, len(num)):
        if num[j] == '1':
            new += 2**(len(num) - j - 1)
    return new


def comb(num, mas):
    num = f'{num:036b}'
    mas = list(mas)
    num = list(num)
    for i in range(0, len(num)):
        if mas[i] == '0':
            num[i] = '0'
        if mas[i] == '1':
            num[i] = '1'

    return to_int(num)


def execute(arg, val):
    if arg == 'm':
        global mask
        mask = val
    else:
        memory[arg] = comb(int(val), mask)



def comb_address(num, mas):
    mas = list(mas)
    m = to_int(mas)

    x = ['0'] * 36
    for i in range(0, len(mas)):
        if mas[i] == 'X':
            x[i] = '1'
    x = to_int(x)

    num = num & ~x
    comb = 0

    for y in range(0, 2**mask.count('X')):
        yield comb | num | m
        comb = ((comb | ~x) + 1) & x


def execute_v2(arg, val):
    if arg == 'm':
        global mask
        mask = val
    else:
        for ad in comb_address(int(arg), mask):
            memory2[ad] = int(val)


with open(r"C:\Users\Maya\Desktop\Uni\WS_2021\AoC\input14") as f:
    lines = f.readlines()

    dec = [decode_line(l) for l in lines]

    memory = {}

    mask = 0

    for d in dec:
        execute(*d)

    print(sum(memory.values()))

    memory2 = {}

    mask = 0

    for d in dec:
        execute_v2(*d)

    print(sum(memory2.values()))