import shared
import re
from pprint import pprint

pattern = r"Card\s+(\d+): ([\d\s]+) \| ([\d\s]+)"


def load_cards(rows):
    cards = {}
    for row in rows:
        match = re.match(pattern, row)
        card_number = int(match.group(1))
        winning_numbers = set(list(map(int, match.group(2).split())))
        have_numbers = set(list(map(int, match.group(3).split())))
        have = winning_numbers.intersection(have_numbers)
        cards[card_number] = dict(
            winning_numbers=winning_numbers,
            have_numbers=have_numbers,
            have=have,
            have_count=len(have),
            points=0,
            refs=0,
        )
        if have:
            points = 1
            for _ in range(len(have) - 1):
                points *= 2
            cards[card_number]["points"] = points
    return cards


# @shared.profile
def part1(cards):
    print(sum(card["points"] for _, card in cards.items()))


# @shared.profile
def part2(cards):
    card_ids = list(range(len(cards) + 1))
    for card_id, card in cards.items():
        nn = [x + 1 + card_id for x in range(card["have_count"])]
        if not nn:
            continue

        # One point for initially having the card
        for idx in nn:
            cards[idx]["refs"] += 1
        # X more points for how many times this card is referenced
        for _ in range(card["refs"]):
            for idx in nn:
                cards[idx]["refs"] += 1
    # Sum the ref counts, and then add the total cards in
    print(sum(card["refs"] for _, card in cards.items()) + len(cards))


def main():
    rows = [row for row in shared.load_rows(4)]
    cards = load_cards(rows)
    with shared.elapsed_timer() as elapsed:
        part1(cards)
        print("🕒", elapsed())

    with shared.elapsed_timer() as elapsed:
        part2(cards)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()