#!/usr/bin/python3

import re

contained_by={}
contains={}

for rule in open("input.txt", "r"):
    rule=rule.rstrip()
    print("Rule: {}".format(rule))
    m=re.match("^(.*?) bag(s|) contain (.*)$", rule)
    if m:
        bag=m.group(1)
        contain=m.group(3).split(", ")
        for contained in contain:
            m2=re.match('^([0-9]*) (.*?) bag(s|).*$', contained)
            if m2 is not None:
                bag_count=int(m2.group(1))
                bag_name=m2.group(2)
                if not bag in contains:
                    contains[bag]=[]
                contains[bag].append((bag_count, bag_name))
                if not bag_name in contained_by:
                    contained_by[bag_name] = []
                if not bag in contained_by[bag_name]:
                    contained_by[bag_name].append(bag)

looking_for="shiny gold"

bags=[]
for bag in contained_by[looking_for]:
    if bag not in bags:
        bags.append(bag)

print(bags)

new_bags=True
new_bag_list=[b for b in bags]
new_new_bag_list=[]
while new_bags:
    for bag1 in new_bag_list:
        if bag1 not in contained_by:
            continue
        for bag2 in contained_by[bag1]:
            if bag2 not in bags:
                new_new_bag_list.append(bag2)
                bags.append(bag2)
    if (len(new_new_bag_list) > 0):
        new_bag_list=[b for b in new_new_bag_list]
        new_new_bag_list=[]
        new_bags=True
    else:
        new_bag_list=[]
        new_bags=False

print("Found {} bags that can contain {} bags".format(len(bags), looking_for))

def get_count(bag_name):
    count=0
    if not bag_name in contains:
        return 0
    for (num,bag) in contains[bag_name]:
        count+=(get_count(bag)+1)*num
    return count

# now looking for how many bags in bags we're looking for
print("Contains {} bags".format(get_count(looking_for)))