#!/usr/bin/env python3
import aoc20
import sys
import regex as re
import functools


class Node:
    def __init__(self, name, children=[]):
        self.name = name
        self.children = children

    def __iter__(self):
        yield from (child for child, _ in self.children)


def parse(_in):
    nodes = {}
    for line in _in:
        match = re.match(r"(\w+ \w+) bags contain (no other bags|(((\d+) (\w+ \w+)) bags?(, )?)+)\.", line)
        children = [(node_str, int(amount)) for amount, node_str in zip(match.captures(5), match.captures(6))]
        nodes[match[1]] = Node(match[1], children)
    for node in nodes.values():
        node.children = [(nodes[node_str], amount) for node_str, amount in node.children]
    return nodes


def pt1(_in):
    @functools.cache
    def can_contain(node, target):
        return target in node or any(can_contain(child, target) for child in node)

    nodes = parse(_in)
    return len([node for node in nodes.values() if can_contain(node, nodes["shiny gold"])])


def pt2(_in):
    @functools.cache
    def count_children(node):
        return 1 + sum(child[1] * count_children(child[0]) for child in node.children)

    nodes = parse(_in)
    return count_children(nodes["shiny gold"]) - 1


if __name__ == "__main__":
    input = aoc20.read_input(sys.argv[1:], 7)
    # graph(parse(input))
    print(pt1(input))
    print(pt2(input))