open Printf;;

let rec list_of_lines in_file =
  try
    let line = input_line in_file in
    line :: list_of_lines(in_file)
  with End_of_file ->
    close_in in_file;
    []

let is_valid_position_numpad (x, y) =
  x >= 0 && x < 3 && y >= 0 && y < 4 && not (x = 0 && y = 3)

let is_valid_position_keypad (x, y) =
  x >= 0 && x < 3 && y >= 0 && y < 2 && not (x = 0 && y = 0)

let get_keypad_coords symbol = 
  match symbol with
  | '<' -> (0, 1)
  | '>' -> (2, 1)
  | '^' -> (1, 0)
  | 'v' -> (1, 1)
  | 'A' -> (2, 0)
  | _ -> raise (Invalid_argument "Invalid symbol")
  
let get_numpad_coords symbol = 
  if symbol = 'A' then
    (2, 3)
  else
    let n = (Char.code symbol) - (Char.code '0') in
    if n = 0 then
      (1, 3)
    else
      ((n - 1) mod 3, (9 - n) / 3)

let get_dir_vector dir_symbol =
  match dir_symbol with
  | '<' -> (-1, 0)
  | '>' -> (1, 0)
  | '^' -> (0, -1)
  | 'v' -> (0, 1)
  | _ -> raise (Invalid_argument "Invalid direction")

let allowed_moves prev =
  match prev with
  | Some 'A' | None -> ['<'; '^'; '>'; 'v']
  | Some '<' -> ['<'; '^'; 'v']
  | Some '>' -> ['^'; '>'; 'v']
  | Some '^' -> ['<'; '^'; '>']
  | Some 'v' -> ['<'; '>'; 'v']
  | _ -> raise (Invalid_argument "Invalid move")

module CombinationSearch =
struct
  type t = int * (int * (int * int) * char * char list)
  let compare  a b =
    Stdlib.compare a b
end
module CombinationSearchSet = Set.Make(CombinationSearch)

let rec find_best_combination_dijkstra lower_level_cost_fun is_valid_fun to_pos queue =
  match CombinationSearchSet.min_elt_opt queue with
  | None -> raise Not_found
  | Some (score, (partial_score, pos, prev_dir, history)) ->
    if pos = to_pos then begin
      score
    end
    else
      let queue_removed = CombinationSearchSet.remove (score, (partial_score, pos, prev_dir, history)) queue in
      Some prev_dir
      |> allowed_moves
      |> List.filter_map (fun dir ->
        let vec = get_dir_vector dir in
        let new_pos = (fst pos + fst vec, snd pos + snd vec) in
        let lower_cost = lower_level_cost_fun prev_dir dir in
        let finish_cost = lower_level_cost_fun dir 'A' in
        let new_partial_score = partial_score + lower_cost in
        if is_valid_fun new_pos then
          Some (new_partial_score + finish_cost, (new_partial_score, new_pos, dir, dir :: history))
        else
          None)
      |> List.fold_left (fun acc state -> CombinationSearchSet.add state acc) queue_removed
      |> find_best_combination_dijkstra lower_level_cost_fun is_valid_fun to_pos

let find_best_combination lower_level_cost_fun is_valid_fun get_coords_fun from_symbol to_symbol =
  let coords1 = get_coords_fun from_symbol in
  let coords2 = get_coords_fun to_symbol in
  find_best_combination_dijkstra lower_level_cost_fun is_valid_fun coords2 ([(1, (0, coords1, 'A', []))] |> CombinationSearchSet.of_list)

let get_numerical_value str =
  int_of_string (String.sub str 0 (String.length str - 1))

let rec get_keypad_combination_length cache depth from_symbol to_symbol =
  match Hashtbl.find_opt cache (from_symbol, to_symbol, depth) with
  | Some score -> score
  | None ->
    let best = find_best_combination (get_keypad_combination_length cache (depth - 1))
                  is_valid_position_keypad get_keypad_coords from_symbol to_symbol in
    Hashtbl.add cache (from_symbol, to_symbol, depth) best;
    best

let get_numpad_combination_length cache keypad_depth from_symbol to_symbol = 
  let result = find_best_combination (get_keypad_combination_length cache keypad_depth)
      is_valid_position_numpad get_numpad_coords from_symbol to_symbol in
  result

let get_numpad_str_combination_length cache keypad_depth str =
  str
  |> String.fold_left (fun (cost, prev) c -> (cost + get_numpad_combination_length cache keypad_depth prev c, c)) (0, 'A')
  |> fst

let sum_solutions cache keypad_depth lst =
  lst
  |> List.map (fun str -> (get_numpad_str_combination_length cache keypad_depth str, get_numerical_value str))
  |> List.map (fun (solution, value) -> solution * value)
  |> List.fold_left Int.add 0

let init_cache =
  let cache = Hashtbl.create 256 in
  let key_options = ['<'; '>'; '^'; 'v'; 'A'] in
  key_options |> List.iter (fun c1 ->
    key_options |> List.iter (fun c2 -> Hashtbl.add cache (c1, c2, 0) 1));
  cache

let () =
  let () = Printexc.record_backtrace true in
  let f = open_in "input.txt" in
  let lines = list_of_lines f in
  let cache = init_cache in
  let result = sum_solutions cache 2 lines in
  let result2 = sum_solutions cache 25 lines in
  printf "%d\n%d\n" result result2