from tools.aoc import AOCDay
from typing import Any

ARROW_PATHS = {
    ("A", "<"): "v<<",
    ("A", ">"): "v",
    ("A", "^"): "<",
    ("A", "v"): "<v",
    ("^", "<"): "v<",
    ("^", ">"): "v>",
    ("^", "A"): ">",
    ("<", "A"): ">>^",
    ("<", "^"): ">^",
    ("<", "v"): ">",
    ("v", "A"): "^>",
    ("v", "<"): "<",
    ("v", ">"): ">",
    (">", "A"): "^",
    (">", "^"): "<^",
    (">", "v"): "<",
}

KEYPAD_PATHS = {
    ("A", "0"): "<",
    ("A", "1"): "^<<",
    ("A", "2"): "<^",
    ("A", "3"): "^",
    ("A", "4"): "^^<<",
    ("A", "5"): "<^^",
    ("A", "6"): "^^",
    ("A", "7"): "^^^<<",
    ("A", "8"): "<^^^",
    ("A", "9"): "^^^",
    ("0", "A"): ">",
    ("0", "1"): "^<",
    ("0", "2"): "^",
    ("0", "3"): "^>",
    ("0", "4"): "^^<",
    ("0", "5"): "^^",
    ("0", "6"): ">^^",
    ("0", "7"): "^^^<",
    ("0", "8"): "^^^",
    ("0", "9"): "^^^>",
    ("1", "A"): ">>v",
    ("1", "0"): ">v",
    ("1", "2"): ">",
    ("1", "3"): ">>",
    ("1", "4"): "^",
    ("1", "5"): "^>",
    ("1", "6"): ">>^",
    ("1", "7"): "^^",
    ("1", "8"): ">^^",
    ("1", "9"): ">>^^",
    ("2", "A"): "v>",
    ("2", "0"): "v",
    ("2", "1"): "<",
    ("2", "3"): ">",
    ("2", "4"): "<^",
    ("2", "5"): "^",
    ("2", "6"): "^>",
    ("2", "7"): "<^^",
    ("2", "8"): "^^",
    ("2", "9"): "^^>",
    ("3", "A"): "v",
    ("3", "0"): "<v",
    ("3", "1"): "<<",
    ("3", "2"): "<",
    ("3", "4"): "<<^",
    ("3", "5"): "<^",
    ("3", "6"): "^",
    ("3", "7"): "<<^^",
    ("3", "8"): "<^^",
    ("3", "9"): "^^",
    ("4", "A"): ">>vv",
    ("4", "0"): ">vv",
    ("4", "1"): "v",
    ("4", "2"): "v>",
    ("4", "3"): "v>>",
    ("4", "5"): ">",
    ("4", "6"): ">>",
    ("4", "7"): "^",
    ("4", "8"): "^>",
    ("4", "9"): "^>>",
    ("5", "A"): "vv>",
    ("5", "0"): "vv",
    ("5", "1"): "<v",
    ("5", "2"): "v",
    ("5", "3"): "v>",
    ("5", "4"): "<",
    ("5", "6"): ">",
    ("5", "7"): "<^",
    ("5", "8"): "^",
    ("5", "9"): "^>",
    ("6", "A"): "vv",
    ("6", "0"): "<vv",
    ("6", "1"): "<<v",
    ("6", "2"): "<v",
    ("6", "3"): "v",
    ("6", "4"): "<<",
    ("6", "5"): "<",
    ("6", "7"): "<<^",
    ("6", "8"): "<^",
    ("6", "9"): "^",
    ("7", "A"): ">>vvv",
    ("7", "0"): ">vvv",
    ("7", "1"): "vv",
    ("7", "2"): "vv>",
    ("7", "3"): "vv>>",
    ("7", "4"): "v",
    ("7", "5"): "v>",
    ("7", "6"): "v>>",
    ("7", "8"): ">",
    ("7", "9"): ">>",
    ("8", "A"): "vvv>",
    ("8", "0"): "vvv",
    ("8", "1"): "<vv",
    ("8", "2"): "vv",
    ("8", "3"): "vv>",
    ("8", "4"): "<v",
    ("8", "5"): "v",
    ("8", "6"): "v>",
    ("8", "7"): "<",
    ("8", "9"): ">",
    ("9", "A"): "vvv",
    ("9", "0"): "<vvv",
    ("9", "1"): "<<vv",
    ("9", "2"): "<vv",
    ("9", "3"): "vv",
    ("9", "4"): "<<v",
    ("9", "5"): "<v",
    ("9", "6"): "v",
    ("9", "7"): "<<",
    ("9", "8"): "<",
}


def get_arrow_path(code: str, paths: dict[tuple[str, str], str]) -> str:
    cur_pos = "A"
    path = ""
    for c in code:
        if cur_pos != c:
            path += paths[(cur_pos, c)]

        path += "A"
        cur_pos = c

    return path


class Day(AOCDay):
    inputs = [
        [
            (126384, "input21_test"),
            (224326, "input21_jonathan"),
            (163920, "input21"),
        ],
        [
            (279638326609472, "input21_jonathan"),
            (204040805018350, "input21"),
        ],
    ]

    def get_path_len(self, code: str, num_pads: int = 2) -> int:
        if num_pads == 0:
            return len(code)

        if (code, num_pads) in self.DP:
            return self.DP[code, num_pads]

        path = get_arrow_path(code, ARROW_PATHS)
        p_len = 0
        for sub_path in path.split("A")[:-1]:
            p_len += self.get_path_len(sub_path + "A", num_pads - 1)

        self.DP[code, num_pads] = p_len
        return self.DP[code, num_pads]

    def get_full_len(self, num_keypads: int = 2) -> int:
        ans = 0
        for code in self.getInput():
            key_pad_path = get_arrow_path(code, KEYPAD_PATHS)
            ans += self.get_path_len(key_pad_path, num_keypads) * int(code[:-1])

        return ans

    def part1(self) -> Any:
        return self.get_full_len()

    def part2(self) -> Any:
        return self.get_full_len(25)


if __name__ == "__main__":
    day = Day(2024, 21)
    day.run(verbose=True)