(** [populate_grid grid start] does a depth-first search through [grid] to find
    the route from [start] to the end. *)
let populate_grid grid start =
  let costs = Array.make (Aoc.Grid.length grid) max_int in
  let rec step acc cost = function
    | [] -> acc
    | (x, y) :: t ->
        if
          Aoc.Grid.pos_is_valid grid (x, y)
          && Aoc.Grid.get_by_pos grid (x, y) <> '#'
          && costs.(Aoc.Grid.idx_of_pos grid (x, y)) = max_int
        then begin
          costs.(Aoc.Grid.idx_of_pos grid (x, y)) <- cost;
          if Aoc.Grid.get_by_pos grid (x, y) = 'E' then step acc cost t
          else
            step
              ((x - 1, y) :: (x + 1, y) :: (x, y - 1) :: (x, y + 1) :: acc)
              cost t
        end
        else step acc cost t
  in
  let rec dfs cost lst =
    let next_step = step [] cost lst in
    if next_step = [] then costs else dfs (cost + 1) next_step
  in
  let costs = dfs 0 [ start ] in
  costs

(** [manhattan_distance2 p p'] returns the Manhattan distance between two points
    on a 2-D plane. *)
let manhattan_distance2 (x, y) (x', y') = abs (x - x') + abs (y - y')

(** [within_distance pos distance] returns all points that are at most
    [distance] units away from [pos] when measured using the Manhattan distance.
*)
let within_distance (x, y) distance =
  let rec impl' acc y' x' =
    if manhattan_distance2 (x, y) (x', y') > distance then acc
    else impl' ((x', y') :: acc) y' (x' + 1)
  in
  let rec impl acc y' =
    if y' - y > distance then acc
    else impl (impl' acc y' (x - (distance - abs (y - y')))) (y' + 1)
  in
  impl [] (y - distance)

(** [find_cost min_amount depth_first grid length idx] returns the number of
    cheat routes starting at [idx] which have a saving of at least [min_amount]
    and are no longer than [length]. [depth_first] is the cost map, and [grid]
    is the grid. *)
let find_cost min_amount depth_first grid length idx =
  (* because there is only one route through the grid we can specialize and look
     to see how much we can save by going from [idx] to any other grid position
     within [length] units (manhattan distance).  The saving is the cost of
     gettimg to idx' from idx via the old route - the cost via the new route. *)
  let saving idx' =
    let cost = depth_first.(idx) in
    let cost' = depth_first.(idx') in
    cost' - cost
    - manhattan_distance2
        (Aoc.Grid.pos_of_idx grid idx)
        (Aoc.Grid.pos_of_idx grid idx')
  in
  within_distance (Aoc.Grid.pos_of_idx grid idx) length
  |> List.filter (Aoc.Grid.pos_is_valid grid)
  |> List.map (Aoc.Grid.idx_of_pos grid)
  |> List.filter (fun idx' -> depth_first.(idx') <> max_int)
  |> List.filter (fun idx' -> depth_first.(idx') - depth_first.(idx) >= 0)
  |> List.map saving
  |> List.filter (( <= ) min_amount)
  |> List.length

(** [find_cost_reductions min_amount cheat_length (grid, start)] returns the
    number of cheat-routes that can be found in [grid] starting at [start] that
    save at least [min_amount] moves and are no longer than [cheat_length]
    units. *)
let find_cost_reductions min_amount cheat_length (grid, start) =
  let costs = populate_grid grid start in
  Seq.ints 0
  |> Seq.take (Aoc.Grid.length grid)
  |> Seq.map (find_cost min_amount costs grid cheat_length)
  |> Seq.fold_left ( + ) 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, find_cost_reductions 100 2);
      (string_of_int, find_cost_reductions 100 20);
    ]