(** [next_char ch] returns the next character after [ch]. *)
let next_char ch = Char.chr (1 + Char.code ch)

(** [ten] really isn't the digit 10, but is the character after '9'. *)
let ten = next_char '9'

(** The character '1'. *)
let one = '1'

(** [find_trail grid pos0] returns a list of all end points of trails in [grid]
    starting at [pos0]. [pos0] must point to a valid position that contains a
    '0'. The same endpoint may be returned multiple times if there are multiple
    routes to it. *)
let find_trail grid pos0 =
  assert (Aoc.Grid.get_by_pos grid pos0 = '0');
  let add_pos lst pos digit =
    if Aoc.Grid.pos_is_valid grid pos && Aoc.Grid.get_by_pos grid pos = digit
    then pos :: lst
    else lst
  in
  let add_poses lst (x, y) digit =
    let lst = add_pos lst (x - 1, y) digit in
    let lst = add_pos lst (x + 1, y) digit in
    let lst = add_pos lst (x, y - 1) digit in
    let lst = add_pos lst (x, y + 1) digit in
    lst
  in
  let rec find_next acc digit = function
    | [] -> acc
    | h :: t -> find_next (add_poses acc h digit) digit t
  in
  let rec impl acc digit =
    if digit = ten then acc else impl (find_next [] digit acc) (next_char digit)
  in
  impl [ pos0 ] one

(** [find_trails grid] returns the list of list of end-points of trails starting
    at each position in [grid]. The [n]th element of the returned list
    corresponds to the trails starting at index [n]. *)
let find_trails grid =
  let rec impl acc idx =
    if idx >= Aoc.Grid.length grid then acc
    else if Aoc.Grid.get_by_idx grid idx <> '0' then impl acc (idx + 1)
    else (* grid_get_by_idx grid idx = 0 *)
      impl (find_trail grid (Aoc.Grid.pos_of_idx grid idx) :: acc) (idx + 1)
  in
  impl [] 0 |> List.rev

(** [part sort_fn grid] returns a count of all trails in [grid], before counting
    the trails for each grid index are sorted by [sort_fn]. *)
let part sort_fn grid =
  find_trails grid |> List.map sort_fn |> List.map List.length
  |> List.fold_left ( + ) 0

let _ =
  Aoc.main Aoc.Grid.of_file
    [
      (string_of_int, part (List.sort_uniq Stdlib.compare));
      (string_of_int, part Fun.id);
    ]