from tools.aoc import AOCDay
from tools.math import magnitude
from typing import Any


class Day(AOCDay):
    inputs = [
        [
            (55312, "input11_test"),
            (203228, "input11"),
        ],
        [
            (240884656550923, "input11"),
        ],
    ]

    def iter_stone(self, stone: int, times: int) -> int:
        if times == 0:
            return 1

        if (stone, times) in self.DP:
            return self.DP[(stone, times)]

        if stone == 0:
            result = self.iter_stone(1, times - 1)
        elif magnitude(stone) % 2 == 1:
            div = magnitude(stone) // 2 + 1
            result = self.iter_stone(stone % (10**div), times - 1) + self.iter_stone(stone // (10**div), times - 1)
        else:
            result = self.iter_stone(stone * 2024, times - 1)

        self.DP[(stone, times)] = result
        return result

    def part1(self) -> Any:
        return sum(self.iter_stone(x, 25) for x in map(int, self.getInput().split()))

    def part2(self) -> Any:
        return sum(self.iter_stone(x, 75) for x in map(int, self.getInput().split()))


if __name__ == "__main__":
    day = Day(2024, 11)
    day.run(verbose=True)