(** [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

let visit grid has_visited cost ((x, y) as state) =
  if not (Aoc.Grid.pos_is_valid grid state) then []
  else if has_visited.(Aoc.Grid.idx_of_pos grid state) then []
  else begin
    has_visited.(Aoc.Grid.idx_of_pos grid state) <- true;
    [
      (cost + 1, (x - 1, y));
      (cost + 1, (x + 1, y));
      (cost + 1, (x, y - 1));
      (cost + 1, (x, y + 1));
    ]
  end

let find_cost grid start =
  let has_visited =
    Array.init (Aoc.Grid.length grid) (fun x ->
        Aoc.Grid.get_by_idx grid x = '#')
  in
  match
    dijkstra (visit grid has_visited)
      (fun x -> Aoc.Grid.get_by_pos_opt grid x = Some 'E')
      [ (0, start) ]
  with
  | None -> failwith "find_cost"
  | Some (cost, _) -> cost

module IntMap = Map.Make (Int)

let update_value = function None -> Some 1 | Some x -> Some (x + 1)

let find_cost_reductions grid start =
  let max_cost = find_cost grid start in
  let rec impl acc idx =
    if idx >= Aoc.Grid.length grid then acc
    else if Aoc.Grid.get_by_idx grid idx <> '#' then impl acc (idx + 1)
    else
      let grid' = Aoc.Grid.update_idx grid idx '.' in
      let new_cost = find_cost grid' start in
      impl (IntMap.update (max_cost - new_cost) update_value acc) (idx + 1)
  in
  impl IntMap.empty 0

let print_map map = IntMap.iter (fun k v -> Printf.printf "%d -> %d\n" k v) map

let part1 (grid, start) =
  let map = find_cost_reductions grid start in
  print_map map;
  let map = IntMap.filter (fun k _ -> k >= 100) map in
  IntMap.fold (fun _ v a -> v + a) map 0

let find_start grid =
  match Aoc.Grid.idx_from_opt grid 0 'S' with
  | None -> failwith "find_start"
  | Some x -> Aoc.Grid.pos_of_idx grid x

let data_of_file fname =
  let grid = Aoc.Grid.of_file fname in
  let start = find_start grid in
  (grid, start)

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