from aoc.input import get_input
import time
import collections
import re
from aoc.partselector import part_one, part_two
import functools

def pw(line):
    l =  re.split(" bags?,?(?: contain)? ?", line.strip('.'))
    return list(filter(len, l))

def p1():
    inp = get_input(pw)
    pwd = dict()
    for sample in inp:
        data = []
        for element in sample[1:]:
            if 'no' in element:
                continue
            data.append(tuple(element.split(' ', 1)))
        pwd[sample[0]] = data

    @functools.lru_cache(maxsize=256)
    def contains_gold(y):
        return 'shiny gold' == y or any(map(lambda x: contains_gold(x[1]), pwd[y]))

    s = sum(map(lambda x: any(map(lambda y: contains_gold(y[1]), pwd[x])), pwd))


    print(s)
    return pwd


def p2(r):
    pwd = r
    count = 0

    @functools.lru_cache(maxsize=256)
    def count_bags(y):
        count = 1
        for j in pwd[y]:
            tmpcount = count_bags(j[1])
            count += tmpcount * int(j[0])
        return count

    ret = 0
    for t in pwd['shiny gold']:
        ret += count_bags(t[1]) * int(t[0])

    print(ret)
    return ret

if part_one():
    start = time.time()
    result1 = p1()
    print(round(1000*(time.time() - start), 2), 'ms')


if part_two():
    start = time.time()
    p2(result1)
    print(round(1000*(time.time() - start), 2), 'ms')