#advent of code 2020
#day 14
#first puzzle where I was stuck with wrong answer even tho the code look ok
#trimming down the amount of code helped a lot, since there seemed to be a 
#wrong variable consider during comparison
#probably should've used binary values instead of converting numbers
#to binary manually

import re
import itertools

def GetBit(decimal): #convert decimal to bit as a list
	b = [];
	for p in range(35,-1,-1):
		power = pow(2,p);
		x1 = decimal//power;
		b.append(x1);
		decimal -= x1*power;
	return b;
		
def GetDec(bit): #convert bit (list) to integer
	dec = 0;
	for index, bv in enumerate(bit):
		dec += bv*pow(2,len(bit)-1-index);
	return dec;

def GetAddresses(MemoryAddress, CurrentMask):
	addresses = [];
	BitAddress = GetBit(MemoryAddress);
	Xlocations = [];
	for options in itertools.product([0,1],repeat=CurrentMask.count("X")):
		NewAddress = [];
		xcount = 0;
		for bi, ba in enumerate(BitAddress):
			if CurrentMask[bi] == "X": 
				NewAddress.append(options[xcount]);
				xcount += 1;
			elif CurrentMask[bi] == '1':
				NewAddress.append(1);
			else:
				NewAddress.append(ba);
		addresses.append(GetDec(NewAddress));
	return addresses;
	
mem1 = dict();
mem2 = dict();
mask = None;

PuzzleInput = open("14.in","r");
for line in PuzzleInput:
	line = line[:-1];
	if line[2] == "s": #a bit obfuscated "if it starts with 'mask = ' "
		mask = list(line.replace("mask = ",""));
	else:
		digits = [int(v) for v in re.findall(r'\d+',line)];
		addr, value = digits;
		bval = GetBit(value);
		for i,m in enumerate(mask):
			if m == "X": continue;
			bval[i] = int(m);
		mem1[int(addr)] = GetDec(bval);
		#part 2
		for FloatAddress in GetAddresses(addr,mask):
			mem2[FloatAddress] = value;
PuzzleInput.close();
	
p1 = 0;
p2 = 0;
for a1 in mem1: p1 += mem1[a1];
for a2 in mem2: p2 += mem2[a2];
print("part 1 =",p1);
print("part 2 =",p2);