let pos_of_numeric_grid c =
  match c with
  | '7' -> (0, 0)
  | '8' -> (1, 0)
  | '9' -> (2, 0)
  | '4' -> (0, 1)
  | '5' -> (1, 1)
  | '6' -> (2, 1)
  | '1' -> (0, 2)
  | '2' -> (1, 2)
  | '3' -> (2, 2)
  | '0' -> (1, 3)
  | 'A' -> (2, 3)
  | _ -> raise (invalid_arg "pos_of_numeric_grid")

let pos_of_dir_grid c =
  match c with
  | '^' -> (1, 3)
  | 'A' -> (2, 3)
  | '<' -> (0, 4)
  | 'v' -> (1, 4)
  | '>' -> (2, 4)
  | _ -> raise (invalid_arg "pos_of_dir_grid")

let invalid_x, invalid_y = (0, 3)

let find_shortest (sx, sy) (fx, fy) =
  let b = Buffer.create 6 in
  let result = [] in
  let result =
    if fx <> invalid_x || sy <> invalid_y then begin
      Buffer.reset b;
      if fx > sx then Buffer.add_string b (String.make (fx - sx) '>')
      else if fx < sx then Buffer.add_string b (String.make (sx - fx) '<');
      if fy > sy then Buffer.add_string b (String.make (fy - sy) 'v')
      else if fy < sy then Buffer.add_string b (String.make (sy - fy) '^');
      Buffer.add_char b 'A';
      Buffer.contents b :: result
    end
    else result
  in
  let result =
    if sx <> invalid_x || fy <> invalid_y then begin
      Buffer.reset b;
      if fy > sy then Buffer.add_string b (String.make (fy - sy) 'v')
      else if fy < sy then Buffer.add_string b (String.make (sy - fy) '^');
      if fx > sx then Buffer.add_string b (String.make (fx - sx) '>')
      else if fx < sx then Buffer.add_string b (String.make (sx - fx) '<');
      Buffer.add_char b 'A';
      Buffer.contents b :: result
    end
    else result
  in
  List.sort_uniq compare result

let[@warning "-32"] print_list lst =
  let rec impl = function
    | [] -> ()
    | h :: t ->
        Printf.printf "; %s" h;
        impl t
  in
  match lst with
  | [] -> print_endline "[]"
  | h :: t ->
      Printf.printf "[%s" h;
      impl t;
      Printf.printf "]\n"

module CharPair = struct
  type t = char * char

  let compare (x, y) (x', y') =
    match compare y y' with 0 -> compare x x' | c -> c
end

module CharPairMap = Map.Make (CharPair)

let routes pos_of_grid locs =
  let rec impl acc pos =
    let start = pos_of_grid pos in
    function
    | [] -> acc
    | h :: t ->
        impl
          (CharPairMap.add (pos, h) (find_shortest start (pos_of_grid h)) acc)
          pos t
  in
  let rec impl' acc = function
    | [] -> acc
    | h :: t -> impl' (impl acc h locs) t
  in
  impl' CharPairMap.empty locs

let num_grid = [ '0'; '1'; '2'; '3'; '4'; '5'; '6'; '7'; '8'; '9'; 'A' ]
let num_routes = routes pos_of_numeric_grid num_grid
let dir_grid = [ '<'; '>'; 'v'; '^'; 'A' ]
let dir_routes = routes pos_of_dir_grid dir_grid

let cross_product lst lst' =
  let rec impl' acc app' = function
    | [] -> acc
    | h :: t -> impl' ((h ^ app') :: acc) app' t
  in
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (impl' acc h lst) t
  in
  if List.is_empty lst then lst'
  else if List.is_empty lst' then lst
  else impl [] lst'

let find_route route_map presses =
  let rec impl acc from presses =
    match Seq.uncons presses with
    | None -> acc
    | Some (h, t) ->
        let routes = CharPairMap.find (from, h) route_map in
        impl (cross_product acc routes) h t
  in
  impl [] 'A' (String.to_seq presses)

let find_routes route_map press_list =
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (acc @ find_route route_map h) t
  in
  impl [] press_list

let min_length = List.fold_left (fun acc x -> min acc (String.length x)) max_int

let human_insns2 count code =
  let rec robot_round codes n =
    if n >= count then codes
    else robot_round (find_routes dir_routes codes) (n + 1)
  in
  let first_robot_insns = find_routes num_routes [ code ] in
  let human = robot_round first_robot_insns ~-1 in
  min_length human

let complexity2 count code =
  let len = human_insns2 count code in
  let num = int_of_string (String.sub code 0 (String.length code - 1)) in
  num * len

let part1a count codes =
  List.map (complexity2 count) codes |> List.fold_left ( + ) 0

let initial_cost_map grid =
  let rec impl' acc f = function
    | [] -> acc
    | h :: t -> impl' (CharPairMap.add (f, h) 1 acc) f t
  in
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (impl' acc h grid) t
  in
  impl CharPairMap.empty grid

let calc_costs cost_map steps =
  let rec impl acc from seq =
    match Seq.uncons seq with
    | None -> acc
    | Some (h, t) -> impl (acc + CharPairMap.find (from, h) cost_map) h t
  in
  impl 0 'A' (String.to_seq steps)

let get_next_level_costs route_map cost_map =
  let impl routes =
    List.fold_left (fun acc x -> min acc (calc_costs cost_map x)) max_int routes
  in
  CharPairMap.map impl route_map

let human_insns3 count code =
  let first_round = initial_cost_map dir_grid in
  let rec robot_round cost_map n =
    if n >= count then cost_map
    else robot_round (get_next_level_costs dir_routes cost_map) (n + 1)
  in
  let number_costs = robot_round first_round 0 in
  let number_costs = get_next_level_costs num_routes number_costs in
  let human = calc_costs number_costs code in
  human

let complexity3 count code =
  let len = human_insns3 count code in
  let num = int_of_string (String.sub code 0 (String.length code - 1)) in
  num * len

let part1b count codes =
  List.map (complexity3 count) codes |> List.fold_left ( + ) 0

let _ =
  Aoc.main Aoc.strings_of_file
    [
      (string_of_int, part1a 1);
      (string_of_int, part1b 2);
      (string_of_int, part1b 25);
    ]

(* too high 524110008179112 *)