from aocd.models import Puzzle
from aocd import submit
from dataclasses import dataclass
import math
import re


def get_winning_numbers(input_data: str) -> list:
    """
    Parse input_data, extract all cards and turn them into sets.
    Return the intersection between numbers and winning numbers.
    """

    # separate cards into a list of strings and remove excess whitespace
    pairs = [ " ".join(card.split()) for card in input_data.split("\n") ]
    # extract a set of winning numbers
    winners = [ set(pair.split(":")[1].split("|")[0].strip().split(" ")) for pair in pairs ]
    # extract a set of scratch numbers
    numbers = [ set(pair.split(":")[1].split("|")[1].strip().split(" ")) for pair in pairs ]
    # calculate set intersection between numbers and winning numbers
    intersections = [ numbers[i].intersection(winners[i]) for i in range(len(numbers)) ]

    return intersections

def calculate_score(winning_numbers:list) -> list:
    """
    Calculate the score for each number
    - 0 points for no match
    - 1 point for 1 match
    - double the points for each successive match
    """
    scores = []
    for i, win in enumerate(winning_numbers):
        if len(win) == 0:
            scores.append(0)
        if len(win) == 1:
            scores.append(1)
        if len(win) > 1:
            scores.append(2**(len(win) - 1))
    return scores


if __name__ == "__main__":
    # get puzzle and parse data
    puzzle = Puzzle(year=2023, day=4)

    # parse_input(puzzle.input_data)
    winning_numbers = get_winning_numbers(puzzle.input_data)

    answer_a = sum(calculate_score(winning_numbers))
    print(f"{answer_a}")
    submit(answer_a, part="a", day=4, year=2023)

    # part b:
    answer_b = 0
    print(f"{answer_b}")
    #submit(answer_b, part="b", day=4, year=2023)