(** A Pair of characters *)
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)

(** [pos_of_numeric_grid c] returns the [(x, y)] position of [c] in the numeric
    grid. *)
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")

(** [pos_of_numeric_grid c] returns the [(x, y)] position of [c] in the
    direction grid. *)
let pos_of_dir_grid c =
  (* Implementation note: We chose 'A' to have the same position in both grids
     so that there is only one location for the hole. *)
  match c with
  | '^' -> (1, 3)
  | 'A' -> (2, 3)
  | '<' -> (0, 4)
  | 'v' -> (1, 4)
  | '>' -> (2, 4)
  | _ -> raise (invalid_arg "pos_of_dir_grid")

(** Location of the hole which the robot can not go to. *)
let invalid_x, invalid_y = (0, 3)

(** [find_paths start finish] returns a list of paths (using the direction
    keypad to get from [start] to [finish] positions.

    The routing is picked to avoid the invalid location. *)
let find_paths (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
  result

(** [routes pos_of_grid locs] returns a map of [(start, finish)] pairs mapping
    to a list of paths for getting to that route. [locs] are the locations on
    the grid to investiagte. [pos_of_grid] gives the location of each of the
    [locs]. The returned map contains routes from each element in [locs] to
    every element. *)
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_paths 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

(** Elements on the number grid *)
let num_grid = [ '0'; '1'; '2'; '3'; '4'; '5'; '6'; '7'; '8'; '9'; 'A' ]

(** Routes between positions on the number grid *)
let num_routes = routes pos_of_numeric_grid num_grid

(** Elements on the direction grid *)
let dir_grid = [ '<'; '>'; 'v'; '^'; 'A' ]

(** Routes between positions on the direction grid. *)
let dir_routes = routes pos_of_dir_grid dir_grid

(** [initial_cost_map grid] returns a map for the initial costs (1) of moving
    between different positions on [grid]. *)
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

(** [calc_cost cost_map steps] calculates the cost of following [steps].
    [cost_map] gives the cost of moving between each position. *)
let calc_cost 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)

(** [get_next_level_costs route_map cost_map] gets the cost map which
    corresponds to [route_map] with costs [cost_map]. *)
let get_next_level_costs route_map cost_map =
  let impl routes =
    List.map (calc_cost cost_map) routes
    |> List.fold_left (fun acc x -> min acc x) max_int
  in
  CharPairMap.map impl route_map

(** [min_code_cost count code] returns the number of buttons a human needs to
    press to get [code] entered when indirected through [count] robots. *)
let min_code_cost count code =
  let number_costs =
    Aoc.apply_n count
      (get_next_level_costs dir_routes)
      (initial_cost_map dir_grid)
  in
  let number_costs = get_next_level_costs num_routes number_costs in
  calc_cost number_costs code

(** [get_code_complexity count code] returns the complexity of a given code when
    there are [count] robots involved. *)
let get_code_complexity count code =
  let len = min_code_cost count code in
  let num = int_of_string (String.sub code 0 (String.length code - 1)) in
  num * len

(** [part count codes] returns the puzzle rest when there are [count] robots,
    and you need to enter [codes]. *)
let part count codes =
  List.map (get_code_complexity count) codes |> List.fold_left ( + ) 0

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