from tools.aoc import AOCDay
from typing import Any, Dict, List


def getCaveMap(lines: List[List[str]]) -> Dict[str, List[str]]:
    m = {}
    for l in lines:
        if l[0] in m:
            m[l[0]].append(l[1])
        else:
            m[l[0]] = [l[1]]

        if l[1] in m:
            m[l[1]].append(l[0])
        else:
            m[l[1]] = [l[0]]

    return m


def getPaths(caveMap: Dict[str, List[str]], start: str, visited: List[str], noDouble: bool = True) -> List[List[str]]:
    paths = []
    visited.append(start)
    for thisCave in caveMap[start]:
        if thisCave == 'start':
            continue

        if thisCave == 'end':
            paths.append(['end'])
            continue

        foundDouble = False
        if thisCave.islower() and thisCave in visited:
            if noDouble:
                continue
            foundDouble = True

        sub_paths = getPaths(caveMap, thisCave, visited.copy(), noDouble or foundDouble)
        [paths.append(x) for x in sub_paths if 'end' in x]

    return paths


class Day(AOCDay):
    test_solutions_p1 = [10, 19, 226, 5076]
    test_solutions_p2 = [36, 103, 3509, 145643]

    def part1(self) -> Any:
        caveMap = getCaveMap(self.getInputAsArraySplit("-"))
        return len(getPaths(caveMap, 'start', []))

    def part2(self) -> Any:
        caveMap = getCaveMap(self.getInputAsArraySplit("-"))
        return len(getPaths(caveMap, 'start', [], False))