#!/usr/bin/env python3

import sys
import typing
import functools

input_file = sys.argv[1]

with open(input_file) as fd:
    lines = [line.rstrip() for line in fd.readlines()]

numeric_keypad = ["789", "456", "123", " 0A"]
directional_keypad = [" ^A", "<v>"]

vec = tuple[int, int]
directions = {
    "^": (-1, 0),  # ^ North
    ">": (0, 1),  # > East
    "v": (1, 0),  # v South
    "<": (0, -1),  # < West
}
directional_keypad_buttons = tuple(directions.keys()) + ("A",)

complexity = int(sys.argv[2])
keypads = [numeric_keypad] + ([directional_keypad] * complexity)


def but_pos(but: str, keypad: list[str]) -> vec:
    for i, line in enumerate(keypad):
        if but in line:
            return i, line.index(but)
    raise IndexError(f"No such button: {but} in {keypad}")


def in_bounds(i: int, j: int, keypad: list[str]) -> bool:
    if j not in range(3) or i not in range(len(keypad)):
        return False
    return keypad[i][j] != " "


# FIFTH TRY
# Using 2 as a base


@functools.cache
def press(buts: str, depth: int) -> int:
    if depth == len(keypads):
        return len(buts)
    keypad = keypads[depth]
    i, j = but_pos("A", keypad)
    nums = 0
    for but in buts:
        nnbuts = ""
        ni, nj = but_pos(but, keypad)
        bounded = True
        while nj < j:
            nnbuts += "<"
            j -= 1
            bounded &= in_bounds(i, j, keypad)
        while ni > i:
            nnbuts += "v"
            i += 1
            bounded &= in_bounds(i, j, keypad)
        while ni < i:
            nnbuts += "^"
            i -= 1
            bounded &= in_bounds(i, j, keypad)
        while nj > j:
            nnbuts += ">"
            j += 1
            bounded &= in_bounds(i, j, keypad)
        if not bounded:
            nnbuts = nnbuts[::-1]
        nnbuts += "A"
        nums += press(nnbuts, depth + 1)
    return nums


score = 0
for code in lines:
    print("Code", code)
    topresses = press(code, 0)

    numpart = int("0" + code.replace("A", ""))
    print(f"{topresses=} * {numpart=}")
    score += topresses * numpart

print(score)


# FOURTH TRY


# def press4(buts: str) -> int:
#     poss = [list(but_pos("A")) for keypad in keypads]
#     indexes = [0 for _ in keypads]
#     combis = ["" for _ in keypads]
#     combis[0] = buts
#     while indexes[0] != len(buts):
#         for k in len(keypads):
#         pass
#     return 0
#
#
# score = 0
# for code in lines:
#     print("Code", code)
#
#     topresses = press4(code)
#
#     numpart = int("0" + code.replace("A", ""))
#     print(f"{topresses=} * {numpart=}")
#     score += topresses * numpart
#
# print(score)

sys.exit(0)

# THIRD TRY


def press3(buts: str, depth: int) -> typing.Generator[str, None, None]:
    if depth >= len(keypads):
        yield from buts
        return

    keypad = keypads[::-1][depth]
    but_poss: dict[str, vec] = dict()
    for i, line in enumerate(keypad):
        for j, but in enumerate(line):
            but_poss[but] = i, j
    i, j = but_poss["A"]
    ai, ij = but_poss[" "]

    for but in press3(buts, depth + 1):
        nnbuts = ""
        ni, nj = but_poss[but]
        bounded = True
        if nj < j:
            nnbuts += "<" * (j - nj)
            j = nj
            bounded &= (ai, ij) != (i, j)
        if ni > i:
            nnbuts += "v" * (ni - i)
            i = ni
            bounded &= (ai, ij) != (i, j)
        if ni < i:
            nnbuts += "^" * (i - ni)
            i = ni
            bounded &= (ai, ij) != (i, j)
        if nj > j:
            nnbuts += ">" * (nj - j)
            j = nj
        if not bounded:
            nnbuts = nnbuts[::-1]
        nnbuts += "A"
        yield from nnbuts


score = 0
for code in lines:
    print("Code", code)

    topresses = 0
    for _ in press3(code, 0):
        topresses += 1

    numpart = int("0" + code.replace("A", ""))
    print(f"{topresses=} * {numpart=}")
    score += topresses * numpart

print(score)

# SECOND TRY

# # Shouldn't move over the bounds, repeat movements when possible, also use movements
# # that start further to A then closer to A, because we're going to press A after anyways
# moves = {
#     "AA": "",
#     "A^": "<",
#     "A>": "v",
#     "Av": "<v",  # sort
#     "A<": "v<<",  # hole avoid
#     "^A": ">",
#     "^^": "",
#     "^>": "v>",  # sort
#     "^v": "v",
#     "^<": "v<",  # hole avoid
#     ">A": "^",
#     ">^": "<^",  # sort
#     ">>": "",
#     ">v": "<",
#     "><": "<<",
#     "vA": "^>",  # symetric. but lower layer sort!
#     "v^": "^",
#     "v>": ">",
#     "vv": "",
#     "v<": "<",
#     "<A": ">>^",  # hole avoid
#     "<^": ">^",  # hole avoid
#     "<>": ">>",
#     "<v": ">",
#     "<<": "",
# }
#
#
# def press(buts: str, depth: int) -> str:
#     if depth == len(keypads):
#         return buts
#     keypad = keypads[depth]
#     numerical = keypad == numeric_keypad
#     prev_but = "A"
#     i, j = but_pos(prev_but, keypad)
#     nbuts = ""
#     for but in buts:
#         if numerical:
#             nnbuts = ""
#             ni, nj = but_pos(but, keypad)
#             # right before down
#             # up before left
#             bounded = True
#             while nj < j:
#                 nnbuts += "<"
#                 j -= 1
#                 bounded &= in_bounds(i, j, keypad)
#             while ni > i:
#                 nnbuts += "v"
#                 i += 1
#                 bounded &= in_bounds(i, j, keypad)
#             while ni < i:
#                 nnbuts += "^"
#                 i -= 1
#                 bounded &= in_bounds(i, j, keypad)
#             while nj > j:
#                 nnbuts += ">"
#                 j += 1
#                 bounded &= in_bounds(i, j, keypad)
#             if not bounded:
#                 nnbuts = nnbuts[::-1]
#             nbuts += nnbuts
#         else:
#             move = prev_but + but
#             nbuts += moves[move]
#         nbuts += "A"
#         prev_but = but
#     return press(nbuts, depth + 1)
#
#
# score = 0
# for code in lines:
#     print("Code", code)
#     presses = press(code, 0)
#     print(f"{presses=}")
#     topresses = len(presses)
#
#     numpart = int("0" + code.replace("A", ""))
#     print(f"{topresses=} * {numpart=}")
#     score += topresses * numpart
#
# print(score)

# FIRST TRY

# keys = {
#     "A": "A",
#     "^": "<A" + ">",
#     ">": "VA" + "^",
#     "V": "V<A" + ">^",
#     "<": "V<<A" + ">>^",
# }
#
# def press(but: str, depth: int) -> str:
#     if depth <= 0:
#         return but
#     depth -= 1
#     return "".join(press(b, depth) for b in keys[but])
#
#
# res = press("^", complexity)
# # res += press("A", complexity)
# print(len(res), res)