import re
from pprint import pprint


def load_rules(filename):
	re_rule = re.compile(r"(?P<color>.+?) bags contain (?P<contained>.+)")
	re_contents = re.compile(r"(?P<num_bags>\d+)\s+(?P<color>.+?)\s+bags?")

	rules = {}
	with open(filename, "r", encoding="utf-8") as f:
		for line in f:
			m = re_rule.match(line.rstrip())
			assert m, line
			color = m.group("color")
			contained = m.group("contained")
			rules[color] = {}
			if contained == "no other bags.":
				continue

			m = re_contents.finditer(contained)
			rules[color] = { x.group("color"): int(x.group("num_bags")) for x in m }

	return rules


def what_can_hold(query_color, rules):
	for outer_color, outer_contents in rules.items():
		if not outer_contents:
			continue

		for inner_color, num_bags in outer_contents.items():
			if inner_color == query_color:
				yield outer_color
				yield from what_can_hold(outer_color, rules)


def bag_contains(query_color, rules):
	rule = rules[query_color]
	if not rule:
		return 0

	x = 0
	for color, multiplier in rule.items():
		x += multiplier * (bag_contains(color, rules) + 1)

	return x

rules = load_rules("data.txt")

wch = set(what_can_hold("shiny gold", rules))
print(f"Shiny gold can be held by {len(wch)} bag colors.")

bc = bag_contains("shiny gold", rules)
print(f"A single shiny gold bag requires {bc} individual bags.")