import aoclib

DAY = 15
TEST_SOLUTION_PART1 = 436
TEST_SOLUTION_PART2 = 175594


def part1(test_mode=False, count=2020):
    my_input = aoclib.getInputAsArraySplit(day=DAY, return_type=int, split_char=",", test=test_mode)

    memory_last = {v: i for i, v in enumerate(my_input)}
    memory_prev = {}
    last_spoken = my_input[-1]

    for i in range(len(my_input), count):
        if last_spoken not in memory_prev:
            last_spoken = 0
        else:
            last_spoken = i - 1 - memory_prev[last_spoken]

        if last_spoken in memory_last:
            memory_prev[last_spoken] = memory_last[last_spoken]
        memory_last[last_spoken] = i

    return last_spoken


def part2(test_mode=False):
    return part1(test_mode, count=30000000)