from functools import cache
from collections import defaultdict


ENTRADA_TESTE = "125 17"
ENTRADA_1 = "773 79858 0 71 213357 2937 1 3998391"

# Valores para testes de geração das pedras
TESTES_ENTRADA = [
    [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_para_lista_de_int(entrada):
    return [int(x) for x in entrada.split()]


# Versão simples ou ingênua
# Usa listas para simular o processo de geração de pedras
# Mais fica muito lenta depois de 30 ou 40 lances (não pode ser usada na parte 2)
def simule(entrada: list[int]) -> list[int]:
    saída = []
    for i in entrada:
        if i == 0:
            i = 1
        elif (ls := len(s := str(i))) % 2 == 0:
            ls //= 2
            saída += [int(s[0:ls]), int(s[ls:])]
            continue
        else:
            i *= 2024
        saída.append(i)

    return saída


# Versão ingênua, mas sem usar append.
# Usa uma lista fixa com o dobro do tamanho da entrada
# Mais rápida que a versão ingênua, mas ainda muito lenta para 75 lances
def simule_sem_append(entrada: list[int]) -> list[int]:
    saída = [0] * len(entrada) * 2
    posição = 0
    for i in entrada:
        if i == 0:
            i = 1
        elif (ls := len(s := str(i))) % 2 == 0:
            ls //= 2
            saída[posição] = int(s[0:ls])
            saída[posição + 1] = int(s[ls:])
            posição += 2
            continue
        else:
            i *= 2024
        saída[posição] = i
        posição += 1
    return saída[:posição]


# Soluções para parte 2


# Verifica a pedra e calcula as pedras resultantes
@cache
def multiplique_pedra(pedra: int) -> list[int]:
    if pedra == 0:
        return [1]
    elif (ls := len(s := str(pedra))) % 2 == 0:
        ls //= 2
        return [int(s[0:ls]), int(s[ls:])]
    else:
        return [pedra * 2024]


# Simula o processo de geração de pedras, usando o cache para gerar as pedras
# Não resolve o problema, pois ainda usa listas para simular o processo
def simule_com_cache(entrada: list[int]) -> list[int]:
    saída = [0] * len(entrada) * 2
    posição = 0
    for i in entrada:
        for pedra in multiplique_pedra(i):
            saída[posição] = pedra
            posição += 1
    return saída[:posição]


# Cria um cache com as pedras.
@cache
def multiplique_pedra_lances(pedra: int, lances: int) -> list[int]:
    pedras = [pedra]
    for _ in range(lances):
        pedras = simule_com_cache(pedras)
    return pedras


# Soma as pedras, usando o cache. Faz uma pesquisa em profundidade
# Roda em aproximadamente 27 minutos com 75 lances!
# Funciona, mas longe de ser eficiente
def some_pedras(pedras: list[int], lances: int) -> int:
    PASSO = 25  # Quantidade de lances a simular por chamada
    if lances == 0:
        return len(pedras)
    soma = 0
    for pedra in pedras:
        novas_pedras = multiplique_pedra_lances(pedra, PASSO)
        if lances - PASSO > 0:
            soma += some_pedras(novas_pedras, lances - PASSO)
        else:
            soma += len(novas_pedras)
    return soma


# Melhor solução para parte 2
# Usa um dicionário para armazenar as pedras e suas quantidades
# Roda em aproximadamente 92ms com 75 lances
def soma_pedras_com_dicionario(entrada: list[int], lances: int) -> int:
    d = defaultdict(int)
    for v in entrada:
        d[v] += 1
    for _ in range(lances):
        novo_d = defaultdict(int)
        for pedra in d.keys():
            if pedra == 0:
                novo_d[1] += d[0]
            elif (ls := len(s := str(pedra))) % 2 == 0:
                ls //= 2
                novo_d[int(s[0:ls])] += d[pedra]
                novo_d[int(s[ls:])] += d[pedra]
            else:
                novo_d[pedra * 2024] += d[pedra]
        d = novo_d
    return sum(d.values())


# Segunda melhor solução para parte 2
# Usa uma solução recursiva com cache para calcular as pedras
# Roda em aproximadamente 270ms com 75 lances
@cache
def soluciona_recursivo(pedra, lances) -> int:
    if lances == 0:
        return 1

    if pedra == 0:
        return soluciona_recursivo(1, lances - 1)

    if len(str(pedra)) % 2 == 0:
        sst = str(pedra)
        a, b = sst[: len(sst) // 2], sst[len(sst) // 2 :]
        return soluciona_recursivo(int(a), lances - 1) + soluciona_recursivo(
            int(b), lances - 1
        )

    return soluciona_recursivo(pedra * 2024, lances - 1)


def simule_piscadas(entrada, piscadas) -> int:
    for x in range(piscadas):
        entrada = simule(entrada)
    return len(entrada)


# Parte 1 - Teste
entrada = string_para_lista_de_int(ENTRADA_TESTE)
print("Testes [pedras] número de pedras")
for blink in range(6):
    print(blink + 1, entrada := simule(entrada), len(entrada))
    assert TESTES_ENTRADA[blink] == entrada

entrada = string_para_lista_de_int(ENTRADA_TESTE)
assert simule_piscadas(entrada, 6) == 22

entrada = string_para_lista_de_int(ENTRADA_TESTE)
assert simule_piscadas(entrada, 25) == 55312
print()

# Parte 1 - Entrada 1
entrada_parte_1 = string_para_lista_de_int(ENTRADA_1)
print("Parte 1: ", soma := simule_piscadas(entrada_parte_1, 25))
assert soma == 199982

# Teste da solução recursiva com some_pedras
assert some_pedras(string_para_lista_de_int(ENTRADA_1), 25) == 199982

# A solução com some_pedras funciona para 75 lances, mas demora mais de 27 minutos para rodar!

# Parte 2 - Solução com dicionário +/-92ms
print(
    "Parte 2 com dicionário:",
    soma := soma_pedras_com_dicionario(string_para_lista_de_int(ENTRADA_1), 75),
)
assert soma == 237149922829154
# Parte 2 - Solução recursiva +/- 270ms
print(
    "Parte 2 - recursiva:",
    soma := sum(
        [soluciona_recursivo(s, 75) for s in string_para_lista_de_int(ENTRADA_1)]
    ),
)
assert soma == 237149922829154

# Saída
# Testes [pedras] número de pedras
# 1 [253000, 1, 7] 3
# 2 [253, 0, 2024, 14168] 4
# 3 [512072, 1, 20, 24, 28676032] 5
# 4 [512, 72, 2024, 2, 0, 2, 4, 2867, 6032] 9
# 5 [1036288, 7, 2, 20, 24, 4048, 1, 4048, 8096, 28, 67, 60, 32] 13
# 6 [2097446912, 14168, 4048, 2, 0, 2, 4, 40, 48, 2024, 40, 48, 80, 96, 2, 8, 6, 7, 6, 0, 3, 2] 22

# Parte 1:  199982
# Parte 2 com dicionário: 237149922829154
# Parte 2 - recursiva: 237149922829154