Advent of Code 2024 day 21

from functools import cache
from argparse import ArgumentParser, FileType


def parse(src):
    return [line for line in src.splitlines() if line != '']


def distance(start, end):
    return end[0] - start[0], end[1] - start[1]


numeric_coord = {
    'A': (2, 3),
    '0': (1, 3),
    '1': (0, 2),
    '2': (1, 2),
    '3': (2, 2),
    '4': (0, 1),
    '5': (1, 1),
    '6': (2, 1),
    '7': (0, 0),
    '8': (1, 0),
    '9': (2, 0)
}

coord_numeric = {
    (2, 3): 'A',
    (1, 3): '0',
    (0, 2): '1',
    (1, 2): '2',
    (2, 2): '3',
    (0, 1): '4',
    (1, 1): '5',
    (2, 1): '6',
    (0, 0): '7',
    (1, 0): '8',
    (2, 0): '9'
}


direction_coord = {
    'A': (2, 0),
    '^': (1, 0),
    '>': (2, 1),
    'v': (1, 1),
    '<': (0, 1)
}

coord_direction = {
    (2, 0): 'A',
    (1, 0): '^',
    (2, 1): '>',
    (1, 1): 'v',
    (0, 1): '<'
}


@cache
def move(button, next_button, numeric):
    if button == next_button:
        return 'A'

    if numeric:
        button_coord = numeric_coord
        coord_button = coord_numeric
    else:
        button_coord = direction_coord
        coord_button = coord_direction

    first_coord = button_coord[button]
    next_coord = button_coord[next_button]
    dx, dy = distance(first_coord, next_coord)

    if dx >= 0:
        move_x = '>' * dx
    else:
        move_x = '<' * -dx

    if dy >= 0:
        move_y = 'v' * dy
    else:
        move_y = '^' * -dy

    if dx == 0:
        return move_y + 'A'
    elif dy == 0:
        return move_x + 'A'

    result = list()
    # find the button with the x coord of next_button and the y coord of
    # button, if it exists
    if (next_coord[0], first_coord[1]) in coord_button:
        result.append(move_x + move_y + 'A')

    # similar for y
    if (first_coord[0], next_coord[1]) in coord_button:
        result.append(move_y + move_x + 'A')

    # if we know one will be slower, remove it
    if len(result) != 2 or dx < 0:
        result = result[0]  # '<' is the slowest to get to 'A' from
    else:
        result = result[1]  # 'v' is slower than '^'

    return result


@cache
def get_sequence_length(sequence, indirection, numeric):
    if indirection == 0:
        return len(sequence)

    sequence = 'A' + sequence
    total = 0
    for button, next_button in zip(sequence, sequence[1:]):
        total += get_sequence_length(move(button, next_button, numeric),
                                     indirection - 1, False)
    return total


def type_code(code, indirection):
    length = get_sequence_length(code, indirection + 1, True)
    return length * int(code[:-1])


def main(codes, indirection):
    return sum(type_code(code, indirection) for code in codes)


arg_parser = ArgumentParser()
arg_parser.add_argument('src', type=FileType('r'))
arg_parser.add_argument('indirection', nargs='?', type=int, default=2)

if __name__ == '__main__':
    args = arg_parser.parse_args()
    print(main(parse(args.src.read()), args.indirection))