module IntMap = Map.Make (Int)

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

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

let manhattan_distance (x, y) (x', y') = abs (x - x') + abs (y - y')

let within_distance (x, y) distance =
  let rec impl' acc y' x' =
    if manhattan_distance (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)

let find_cost2 map depth_first grid idx length =
  let saving idx' =
    let cost = depth_first.(idx) in
    let cost' = depth_first.(idx') in
    let saving =
      cost' - cost
      - manhattan_distance
          (Aoc.Grid.pos_of_idx grid idx)
          (Aoc.Grid.pos_of_idx grid idx')
    in
    saving
  in
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (IntMap.update (saving h) update_value acc) t
  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 x -> Aoc.Grid.get_by_idx grid x <> '#')
  |> List.filter (fun x -> depth_first.(x) - depth_first.(idx) >= 0)
  |> impl map

let find_cost_reductions2 grid start cheat_length =
  let depth_first = populate_grid grid start in
  let rec impl acc idx =
    if idx >= Aoc.Grid.length grid then acc
    else
      let acc = find_cost2 acc depth_first grid idx cheat_length in
      impl acc (idx + 1)
  in
  impl IntMap.empty 0

let part2 cheat_length (grid, start) =
  let map = find_cost_reductions2 grid start cheat_length in
  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, part2 2); (string_of_int, part2 20) ]