(** [pairs_of_ints lst] returns a pair from a list of two elements. *)
let pairs_of_ints = function
  | [ h; h' ] -> (h, h')
  | _ -> raise (Invalid_argument "pairs_of_ints")

(** [dijkstra visit check_end states] executes Dijkstra's algorithm.

    [visit cost state] is called to visit [state] with [cost]. It should mark
    [state] as visited, and return a list of [(cost, state)] pairs which contain
    new states to examine. The returned list should be sorted by [cost].

    [check_end state] should return [true] if and only if [state] is an end
    state.

    [states] is a list of [(cost, state)] pairs ordered by [cost].

    [dijkstra] returns [None] if no path is found to the destination. It returns
    [Some (cost, state, remaining_states)] if a route is found. [cost] is the
    cost of getting to [state]. [remaining_states] is a list of the remaining
    states which can be passed back to [dijkstra] if we want to find further
    paths. *)
let rec dijkstra visit check_end =
  let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
  function
  | [] -> None
  | (cost, state) :: t ->
      if check_end state then Some (cost, state)
      else
        let new_states = visit cost state |> List.merge compare_costs t in
        dijkstra visit check_end new_states

type 'a grid = { grid : 'a array; width : int }

(** [grid_is_valid_pos grid (x, y)] returns true if (x, y) is a valid position
*)
let grid_is_valid_pos grid (x, y) =
  x >= 0 && x < grid.width && y >= 0 && y < grid.width

(** Get the index into the grid from an x, y position. *)
let grid_idx_by_pos grid (x, y) = x + (y * grid.width)

(** Set the value of the position (x, y) to v in grid. *)
let grid_set_by_pos grid p v =
  assert (grid_is_valid_pos grid p);
  let idx = grid_idx_by_pos grid p in
  grid.grid.(idx) <- v

(** Get the value of the position (x, y) in grid. *)
let grid_get_by_pos grid p =
  assert (grid_is_valid_pos grid p);
  let idx = grid_idx_by_pos grid p in
  grid.grid.(idx)

(** [grid_of_rocks w rocks] returns a [w * w] grid with [grid.(x + y * w)]
    indicating whether the space is empty ([=max_int]) or which rock it is (0
    based). *)
let grid_of_rocks width rocks =
  let grid = { grid = Array.make (width * width) max_int; width } in
  let add_rock idx p = grid_set_by_pos grid p idx in
  List.iteri add_rock rocks;
  grid

(** [visit grid has_visited count cost pos] visits the location pos marking it
    as visited and returning a list of [(cost, pos)] pairs of next locations to
    examine.

    [grid] is the grid of rocks, [has_visited] is an array of bools indicating
    whether a position has already been visited, and [count] is how many rocks
    have fallen. *)
let visit grid has_visited count cost state =
  if not (grid_is_valid_pos grid state) then []
  else if has_visited.(grid_idx_by_pos grid state) then []
  else if grid_get_by_pos grid state < count then []
  else
    let x, y = state in
    has_visited.(grid_idx_by_pos grid state) <- true;
    [
      (cost + 1, (x + 1, y));
      (cost + 1, (x - 1, y));
      (cost + 1, (x, y + 1));
      (cost + 1, (x, y - 1));
    ]

(** [grid_of_file w fname] returns a grid of width & height [w] populated with
    rocks described in the file [fname]. *)
let grid_of_file width fname =
  Aoc.strings_of_file fname
  |> List.map (Aoc.ints_of_string ~sep:",")
  |> List.map pairs_of_ints |> grid_of_rocks width

(** [find_route_length count grid] calculates the route from the top-left
    position in [grid] to the bottom right if [count] rocks have fallen. It
    returns [None] if no route is possible or [Some (cost, pos)] if the route is
    possible. *)
let find_route_length count grid =
  let has_visited = Array.make (Array.length grid.grid) false in
  dijkstra
    (visit grid has_visited count)
    (( = ) (grid.width - 1, grid.width - 1))
    [ (0, (0, 0)) ]

(** [part1 count rocks] returns how long it takes to navigate the grid [rocks]
    when [count] rocks have fallen. *)
let part1 count rocks =
  match find_route_length count rocks with
  | None -> failwith "part1"
  | Some (cost, _) -> string_of_int cost

(** [part2 start_count grid] returns the location of the first rock to fall into
    [grid] which makes it impossible to get from the top-left to bottom-right.
*)
let part2 width start_count grid =
  (* Implementation notes:
  
     We do this by binary search in impl.  The left_count is a known count of
     rocks that is passable, right_count is a known count that is impassable.

     Once left_count + 1 = right_count we know that right_count is the first
     rock to fall that causes the route to be blocked.

     count_rocks is used to find the number of rocks (and so give an initial
     right_count).
  *)
  let rec count_rocks acc idx =
    if idx >= Array.length grid.grid then acc
    else if grid.grid.(idx) = max_int then count_rocks acc (idx + 1)
    else count_rocks (max acc grid.grid.(idx)) (idx + 1)
  in
  let rec impl left_count right_count =
    if right_count - left_count = 1 then right_count
    else
      let count = (left_count + right_count) / 2 in
      match find_route_length count grid with
      | None -> impl left_count count
      | Some _ -> impl count right_count
  in
  let count = impl start_count (1 + count_rocks 0 0) in
  match Array.find_index (( = ) (count - 1)) grid.grid with
  | None -> failwith "part2"
  | Some idx -> Printf.sprintf "%d,%d" (idx mod width) (idx / width)

(** Width of grid *)
let width = 71

let _ =
  Aoc.main (grid_of_file width)
    [ (Fun.id, part1 1024); (Fun.id, part2 width 1024) ]