from aoc.input import get_input
import copy
import itertools
import time
import collections
import re
from aoc.partselector import part_one, part_two
import functools

def pw(line):
    return line.strip()

def p1():
    inp = get_input(pw)
    mask = ''
    mv = 0
    memory = dict()
    def applymask(mask, v):
        mv = 0
        r = 1
        shift = 0
        for x in reversed(mask):
            if x == '0':
                pass
            if x == '1':
                mv += r
            if x == 'X':
                if (v >> shift) & 0x01 > 0:
                    mv += r
            r *= 2
            shift += 1
        return mv

    for sample in inp:
        t, v = sample.split (' = ')
        if t == 'mask':
            mask = v.strip()
        else:
            cell = t[4:-1]
            memory[cell] = applymask(mask, int(v))

    print(sum(memory.values()))


    return inp


def p2(segments):
    inp = get_input(pw)
    mask = ''
    mv = 0
    memory = dict()
    def applymask(mask, v):
        mv = 0
        r = 1
        shift = 0
        floating = []
        for i, x in enumerate(reversed(mask)):
            if x == '0':
                if (v >> shift) & 0x01 > 0:
                    mv += r
            if x == '1':
                mv += r
            if x == 'X':
                floating.append(i)
            r *= 2
            shift += 1
        for l in range(1, len(floating)+1):
            n = mv
            for f in itertools.combinations(floating, l):
                n = mv
                for f1 in f:
                    n = n | pow(2, f1)
                yield n
        yield mv

    for sample in inp:
        t, v = sample.split (' = ')
        if t == 'mask':
            mask = v.strip()
        else:
            cell = t[4:-1]
            for x in applymask(mask, int(cell)):
                memory[x] = int(v)

    print(sum(memory.values()))


    return inp
    return 0

result1 = None
if part_one():
    start = time.time()
    result1 = p1()
    print(round(1000*(time.time() - start), 2), 'ms')


if part_two():
    start = time.time()
    p2(result1)
    print(round(1000*(time.time() - start), 2), 'ms')