from functools import cache
from collections import defaultdict


TEST_INPUT = "125 17"
INPUT_1 = "773 79858 0 71 213357 2937 1 3998391"

# Values for stone generation tests
TEST_INPUTS = [
    [253000, 1, 7],
    [253, 0, 2024, 14168],
    [512072, 1, 20, 24, 28676032],
    [512, 72, 2024, 2, 0, 2, 4, 2867, 6032],
    [1036288, 7, 2, 20, 24, 4048, 1, 4048, 8096, 28, 67, 60, 32],
    [2097446912, 14168, 4048, 2, 0, 2, 4, 40, 48, 2024, 40, 48, 80, 96, 2, 8, 6, 7, 6, 0, 3, 2],
]


def string_to_int_list(input):
    return [int(x) for x in input.split()]


# Simple or naive version
# Uses lists to simulate the stone generation process
# But gets very slow after 30 or 40 throws (can't be used in part 2)
def simulate(input: list[int]) -> list[int]:
    output = []
    for i in input:
        if i == 0:
            i = 1
        elif (ls := len(s := str(i))) % 2 == 0:
            ls //= 2
            output += [int(s[0:ls]), int(s[ls:])]
            continue
        else:
            i *= 2024
        output.append(i)

    return output


# Naive version, but without using append.
# Uses a fixed list with double the size of the input
# Faster than the naive version, but still too slow for 75 throws
def simulate_without_append(input: list[int]) -> list[int]:
    output = [0] * len(input) * 2
    position = 0
    for i in input:
        if i == 0:
            i = 1
        elif (ls := len(s := str(i))) % 2 == 0:
            ls //= 2
            output[position] = int(s[0:ls])
            output[position + 1] = int(s[ls:])
            position += 2
            continue
        else:
            i *= 2024
        output[position] = i
        position += 1
    return output[:position]


# Solutions for part 2


# Checks the stone and calculates resulting stones
@cache
def multiply_stone(stone: int) -> list[int]:
    if stone == 0:
        return [1]
    elif (ls := len(s := str(stone))) % 2 == 0:
        ls //= 2
        return [int(s[0:ls]), int(s[ls:])]
    else:
        return [stone * 2024]


# Simulates the stone generation process, using cache to generate stones
# Doesn't solve the problem, as it still uses lists to simulate the process
def simulate_with_cache(input: list[int]) -> list[int]:
    output = [0] * len(input) * 2
    position = 0
    for i in input:
        for stone in multiply_stone(i):
            output[position] = stone
            position += 1
    return output[:position]


# Creates a cache with stones.
@cache
def multiply_stone_throws(stone: int, throws: int) -> list[int]:
    stones = [stone]
    for _ in range(throws):
        stones = simulate_with_cache(stones)
    return stones


# Sums stones using cache. Does a depth-first search
# Runs in approximately 27 minutes with 75 throws!
# Works, but far from efficient
def sum_stones(stones: list[int], throws: int) -> int:
    STEP = 25  # Number of throws to simulate per call
    if throws == 0:
        return len(stones)
    total = 0
    for stone in stones:
        new_stones = multiply_stone_throws(stone, STEP)
        if throws - STEP > 0:
            total += sum_stones(new_stones, throws - STEP)
        else:
            total += len(new_stones)
    return total


# Best solution for part 2
# Uses a dictionary to store stones and their quantities
# Runs in approximately 92ms with 75 throws
def sum_stones_with_dictionary(input: list[int], throws: int) -> int:
    d = defaultdict(int)
    for v in input:
        d[v] += 1
    for _ in range(throws):
        new_d = defaultdict(int)
        for stone in d.keys():
            if stone == 0:
                new_d[1] += d[0]
            elif (ls := len(s := str(stone))) % 2 == 0:
                ls //= 2
                new_d[int(s[0:ls])] += d[stone]
                new_d[int(s[ls:])] += d[stone]
            else:
                new_d[stone * 2024] += d[stone]
        d = new_d
    return sum(d.values())


# Second best solution for part 2
# Uses a recursive solution with cache to calculate stones
# Runs in approximately 270ms with 75 throws
@cache
def solve_recursive(stone, throws) -> int:
    if throws == 0:
        return 1

    if stone == 0:
        return solve_recursive(1, throws - 1)

    if len(str(stone)) % 2 == 0:
        sst = str(stone)
        a, b = sst[: len(sst) // 2], sst[len(sst) // 2 :]
        return solve_recursive(int(a), throws - 1) + solve_recursive(
            int(b), throws - 1
        )

    return solve_recursive(stone * 2024, throws - 1)


def simulate_blinks(input, blinks) -> int:
    for x in range(blinks):
        input = simulate(input)
    return len(input)


# Part 1 - Test
input = string_to_int_list(TEST_INPUT)
print("Tests [stones] number of stones")
for blink in range(6):
    print(blink + 1, input := simulate(input), len(input))
    assert TEST_INPUTS[blink] == input

input = string_to_int_list(TEST_INPUT)
assert simulate_blinks(input, 6) == 22

input = string_to_int_list(TEST_INPUT)
assert simulate_blinks(input, 25) == 55312
print()

# Part 1 - Input 1
input_part_1 = string_to_int_list(INPUT_1)
print("Part 1: ", total := simulate_blinks(input_part_1, 25))
assert total == 199982

# Test of recursive solution with sum_stones
assert sum_stones(string_to_int_list(INPUT_1), 25) == 199982

# The solution with sum_stones works for 75 throws, but takes over 27 minutes to run!

# Part 2 - Dictionary solution +/-92ms
print(
    "Part 2 with dictionary:",
    total := sum_stones_with_dictionary(string_to_int_list(INPUT_1), 75),
)
assert total == 237149922829154
# Part 2 - Recursive solution +/- 270ms
print(
    "Part 2 - recursive:",
    total := sum(
        [solve_recursive(s, 75) for s in string_to_int_list(INPUT_1)]
    ),
)
assert total == 237149922829154