import aoclib

DAY = 14
TEST_SOLUTION_PART1 = 51
TEST_SOLUTION_PART2 = 208


def getMaskDict(mask):
    mask_bits = {}
    for i, x in enumerate(mask):
        if x == 'X':
            continue
        mask_bits[pow(2, (len(mask) - i - 1))] = int(x)

    return mask_bits


def part1(test_mode=False):
    my_input = aoclib.getInputAsArraySplit(day=DAY, split_char=" = ", test=test_mode)
    memory = {}
    mask_bits = {}

    for address, value in my_input:
        if address == 'mask':
            mask_bits = getMaskDict(value)
            continue

        address = address.replace('mem[', '').replace(']', '')
        value = int(value)
        for bit in mask_bits:
            if not value & bit and mask_bits[bit]:
                value += bit
            elif value & bit and not mask_bits[bit]:
                value -= bit

        memory[address] = value

    return sum(memory.values())


def part2(test_mode=False):
    my_input = aoclib.getInputAsArraySplit(day=DAY, split_char=" = ", test=test_mode)
    mask = ''
    memory = {}

    for address, value in my_input:
        if address == 'mask':
            mask = value
            continue

        address = format(int(address.replace('mem[', '').replace(']', '')), "036b")
        value = int(value)
        floating_address = ''
        x_index = []
        for i, c in enumerate(address):
            if mask[i] == '0':
                floating_address += address[i]
            else:
                floating_address += mask[i]
                if floating_address[i] == 'X':
                    x_index.append(i)

        x_count = len(x_index)
        if x_count == 0:
            memory[int(floating_address, 2)] = value
        else:
            for x in range(pow(2, x_count)):
                real_address = [c for c in floating_address]
                this_replace = format(x, "0%db" % x_count)

                for i, c in enumerate(this_replace):
                    real_address[x_index[i]] = c

                memory[int("".join(real_address), 2)] = value

    return sum(memory.values())