import aoclib
from functools import lru_cache

DAY = 10
TEST_SOLUTION_PART1 = 220
TEST_SOLUTION_PART2 = 19208


def part1(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, return_type=int, test=test_mode)
    adapters = set(my_input)

    cur_jolt = 0
    count_1 = 0
    count_3 = 0
    for adapter in adapters:
        if adapter - cur_jolt == 1:
            count_1 += 1
        elif adapter - cur_jolt == 3:
            count_3 += 1

        cur_jolt = adapter

    return count_1 * (count_3 + 1)


@lru_cache(None)
def getArrangeCount(adapter_list, cur_jolt):
    count = 0
    found = 0
    if cur_jolt + 1 in adapter_list:
        count += getArrangeCount(adapter_list, cur_jolt + 1)
        found += 1

    if cur_jolt + 2 in adapter_list:
        count += getArrangeCount(adapter_list, cur_jolt + 2)
        found += 1

    if cur_jolt + 3 in adapter_list:
        count += getArrangeCount(adapter_list, cur_jolt + 3)
        found += 1

    if found == 0:
        found = 1

    return count + (found - 1)


def part2(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, return_type=int, test=test_mode)
    adapters = frozenset(my_input)

    return getArrangeCount(adapters, 0) + 1