import re

re_mask = re.compile(r"mask\s+=\s+(?P<mask>.+)")
re_mem  = re.compile(r"mem\[(?P<address>\d+)\]\s+=\s+(?P<value>\d+)")


def part1(filename):
	mask_0 = mask_1 = None
	memory = {}
	with open(filename, "r", encoding="utf-8") as f:
		for line in f:
			s = line.rstrip()
			m = re_mask.match(s)
			if m:
				mask = m.group(1)
				mask_0 = int("".join(["0" if x == "X" else x for x in mask]), 2)
				mask_1 = int("".join(["1" if x == "X" else x for x in mask]), 2)
				continue

			assert mask_0 is not None
			assert mask_1 is not None
			m = re_mem.match(s)
			if m:
				address = int(m.group("address"))
				value = int(m.group("value"))
				memory[address] = value & mask_1 | mask_0
				continue

			raise Exception(f"Parse error: {line}")

		return sum(memory.values())


def part2(filename):
	mask_f = mask_1 = None
	memory = {}
	with open(filename, "r", encoding="utf-8") as f:
		for line in f:
			s = line.rstrip()
			m = re_mask.match(s)
			if m:
				mask = m.group(1)
				mask_f = int("".join(["1" if x == "X" else "0" for x in mask]), 2)
				mask_1 = int("".join(["0" if x == "X" else x for x in mask]), 2)

				f = list(mask)
				f.reverse()
				mask_f = []
				f0 = 0
				for x in [2**i for i, x in enumerate(f) if x == "X"]:
					f0 |= x
					mask_f.append(f0)
				continue

			assert mask_f is not None
			assert mask_1 is not None
			m = re_mem.match(s)
			if m:
				address = int(m.group("address"))
				value = int(m.group("value"))

				memory[address | mask_1] = value
				for x in mask_f:
					memory[address | x] = value

				continue

			raise Exception(f"Parse error: {line}")

		r = sum(memory.values())
		print(r)
		return r


assert part1("test.txt") == 165
r = part1("data.txt")
print(f"1: Answer is {r}")

assert part2("test.txt") == 208
r = part2("data.txt")
print(f"2: Answer is {r}")