zerosleeps

from pathlib import Path
import re
import unittest

def get_raw_input():
    return (Path(__file__).parent/'day_07_input.txt').read_text()

def get_bag_from_input(input):
    return input.strip().split(' bags contain ')[0]

def get_contents_from_input(input):
    if re.search(r'no other bags', input):
        return {}
    else:
        return {
            ' '.join(bag.split()[1:-1]): int(bag.split()[0])
            for bag
            in input.strip().split(' bags contain ')[1].rstrip('.').split(', ')
        }

def parents_for_target(target, rules):
    return [bag for bag, contents in rules.items() if target in contents]

def number_of_children(bag, all_rules):
    count = 0
    if len(all_rules[bag]) > 0:
        for b, q in all_rules[bag].items():
            count += ( q * number_of_children(b, all_rules))
    return count + 1 # Plus one to account for current bag as well

def part_one(raw_input):
    rules = {
        get_bag_from_input(line):get_contents_from_input(line)
        for line
        in raw_input.strip().splitlines()
    }

    candidates = set(parents_for_target('shiny gold', rules))

    growth = True
    while growth:
        new_candidates = { b for c in candidates for b in parents_for_target(c, rules) }
        if new_candidates <= candidates:
            growth = False
        candidates.update(new_candidates)

    return len(candidates)

def part_two(raw_input):
    all_rules = {
        get_bag_from_input(line):get_contents_from_input(line)
        for line
        in raw_input.strip().splitlines()
    }

    return number_of_children('shiny gold',all_rules) - 1

class TestExamples(unittest.TestCase):
    def setUp(self):
        self.example_input = """light red bags contain 1 bright white bag, 2 muted yellow bags.
            dark orange bags contain 3 bright white bags, 4 muted yellow bags.
            bright white bags contain 1 shiny gold bag.
            muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
            shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
            dark olive bags contain 3 faded blue bags, 4 dotted black bags.
            vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
            faded blue bags contain no other bags.
            dotted black bags contain no other bags.
        """

    def test_part_one(self):
        self.assertEqual(part_one(self.example_input), 4)

    def test_part_two_example_one(self):
        self.assertEqual(part_two(self.example_input), 32)

    def test_part_two_example_two(self):
        example_input = """shiny gold bags contain 2 dark red bags.
            dark red bags contain 2 dark orange bags.
            dark orange bags contain 2 dark yellow bags.
            dark yellow bags contain 2 dark green bags.
            dark green bags contain 2 dark blue bags.
            dark blue bags contain 2 dark violet bags.
            dark violet bags contain no other bags.
        """
        self.assertEqual(part_two(example_input), 126)

if __name__ == '__main__':
    print(f'Part one: {part_one(get_raw_input())}')
    print(f'Part two: {part_two(get_raw_input())}')