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_scores(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


def new_rules(input_data: list) -> list:
    """
    Calculate the score for the updated rule set.
    - If your card has n matches, get the next n cards
    - repeat until there are no more matches
    """
    data = input_data.split("\n")
    score, win, total = 0, 0, 0
    cardcounts = []

    for i in range(len(data)):
        cardcounts.append(1)
    for x, row in enumerate(data):
        wins = row.split("|")[0].split(":")[1].split()
        nums = row.split("|")[1].split()
        for num in nums:
            if num in wins:
                win += 1
                if score == 0:
                    score = 1
                else:
                    score *= 2
        for i in range(win):
            cardcounts[x + i + 1] += cardcounts[x]
        total += score
        score, win = 0, 0

    return cardcounts


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_scores(winning_numbers))
    print(f"{answer_a}")
    submit(answer_a, part="a", day=4, year=2023)

    # part b:
    answer_b = sum(new_rules(puzzle.input_data))
    print(f"{answer_b}")
    submit(answer_b, part="b", day=4, year=2023)