let pairs_of_ints = function
  | [ h; h' ] -> (h, h')
  | _ -> raise (Invalid_argument "pairs_of_ints")

let pairs_of_file fname =
  Aoc.strings_of_file fname
  |> List.map (Aoc.ints_of_string ~sep:",")
  |> List.map 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
  | [] -> failwith "dijkstra"
  | (cost, state) :: t ->
      if check_end state then (cost, state)
      else
        let new_states = visit cost state |> List.merge compare_costs t in
        dijkstra visit check_end new_states

type grid = { grid : bool array; width : int }

let grid_is_valid_pos grid (x, y) =
  x >= 0 && x < grid.width && y >= 0 && y < grid.width

let grid_idx_by_pos grid (x, y) = x + (y * grid.width)

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

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)

let make_grid width rocks =
  let grid = { grid = Array.make (width * width) false; width } in
  let rec impl = function
    | [] -> grid
    | p :: t ->
        grid_set_by_pos grid p true;
        impl t
  in
  impl rocks

let visit grid has_visited 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 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));
    ]

let check_end dest state = dest = state

let part1 lst =
  let width = 71 in
  let count = 1024 in
  let grid =
    List.to_seq lst |> Seq.take count |> List.of_seq |> make_grid width
  in
  let has_visited = Array.make (width * width) false in
  let cost, _ =
    dijkstra (visit grid has_visited)
      (check_end (width - 1, width - 1))
      [ (0, (0, 0)) ]
  in
  cost

let _ = Aoc.main pairs_of_file [ (string_of_int, part1) ]