from tools.aoc import AOCDay
from tools.tools import cache
from typing import Any


# 1 way to roll a 3, 3 ways to roll a 4, ...
ROLLS = [(1, 3), (3, 4), (6, 5), (7, 6), (6, 7), (3, 8), (1, 9)]


@cache
def roll(pos1: int, pos2: int, score1: int = 0, score2: int = 0) -> (int, int):
    if score2 >= 21:
        return 0, 1

    p1wins = p2wins = 0

    for sub_universes, pos_inc in ROLLS:
        new_pos1 = 1 + (pos_inc + pos1 - 1) % 10
        # notice the switch-a-roo!
        sub_p2wins, sub_p1wins = roll(pos2, new_pos1, score2, score1 + new_pos1)
        p1wins += sub_p1wins * sub_universes
        p2wins += sub_p2wins * sub_universes

    return p1wins, p2wins


class Day(AOCDay):
    test_solutions_p1 = [739785, 551901]
    test_solutions_p2 = [444356092776315, 272847859601291]

    def part1(self) -> Any:
        pos = list(map(int, [x.split(": ")[1] for x in self.getInput()]))
        scores = [0, 0]
        target_score = 1000
        rolls = 0
        dice_roll = 0
        while scores[0] < target_score and scores[1] < target_score:
            for player in range(2):
                this_roll = 0
                rolls += 3
                for _ in range(3):
                    dice_roll = 1 + (dice_roll % 100)
                    this_roll += dice_roll

                pos[player] = 1 + (this_roll + pos[player] - 1) % 10
                scores[player] += pos[player]
                if scores[player] >= target_score:
                    break

        return min(scores) * rolls

    def part2(self) -> Any:
        pos1, pos2 = map(int, [x.split(": ")[1] for x in self.getInput()])
        wins1, wins2 = roll(pos1, pos2)
        return max(wins1, wins2)