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


def getCaveMap(lines: List[List[str]]) -> Dict[str, List[str]]:
    m = defaultdict(list)
    for l, r in lines:
        m[l].append(r)
        m[r].append(l)

    return m


def getPaths(caveMap: Dict[str, List[str]], start: str, visited: set, noDouble: bool = True) -> int:
    pathCount = 0
    visited.add(start)
    for thisCave in caveMap[start]:
        if thisCave == 'start':
            continue

        if thisCave == 'end':
            pathCount += 1
            continue

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

            foundDouble = True

        pathCount += getPaths(caveMap, thisCave, visited.copy(), noDouble or foundDouble)

    return pathCount


class Day(AOCDay):
    inputs = [
        [
            (10, "test_input12"),
            (19, "test_input12_2"),
            (226, "test_input12_3"),
            (5076, "input12")
        ],
        [
            (36, "test_input12"),
            (103, "test_input12_2"),
            (3509, "test_input12_3"),
            (145643, "input12")
        ]
    ]

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

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


if __name__ == '__main__':
    day = Day(2021, 12)
    day.run(verbose=True)