from aocd.models import Puzzle
from aocd import submit


def valid_games(games: list) -> list:
    """
    Solve part a.
    - check which of the games are possible with the follwing limitations:
        - 12 red
        - 13 green
        - 14 blue
    - add up their IDs
    """
    rules = {"red": 12, "green": 13, "blue": 14}

    valid_ids = []
    for i, game in enumerate(games.split("\n"), start=1):
        throws = game.split(":")[1].split(";")
        valid = True
        for throw in throws:
            cubes = throw.split(",")
            for cube in cubes:
                _, amount, color = cube.split(" ")
                if int(amount) > rules[color]:
                    valid = False
        if valid:
            valid_ids.append(i)
    return valid_ids


def minimum_sets(games: list) -> list:
    """
    Solve part b.
    - figure out the minimum amount of cubes necessary of each color for each throw
    - multiply the values of each minimum number together
    - to solve sum these powers up
    """
    # find all minimas
    minimas = []
    for i, game in enumerate(games.split("\n"), start=1):
        minimum = {"red": 0, "green": 0, "blue": 0}
        throws = game.split(":")[1].split(";")
        for throw in throws:
            cubes = throw.split(",")
            for cube in cubes:
                _, amount, color = cube.split(" ")
                if int(amount) > minimum[color]:
                    minimum[color] = int(amount)
        minimas.append(minimum)

    # calculate all powers
    powers = []
    for minimum in minimas:
        powers.append(minimum["red"] * minimum["green"] * minimum["blue"])
    return powers


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

    # part a:
    answer_a = sum(valid_games(puzzle.input_data))
    print(f"{answer_a}")
    submit(answer_a, part="a", day=2, year=2023)

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