open Printf;;

let rec list_of_lines in_file =
  try
    let line = input_line in_file in
    line :: list_of_lines(in_file)
  with End_of_file ->
    close_in in_file;
    []

let lines_to_2d_array lines =
  let size_x = String.length @@ List.hd lines in
  let size_y = List.length lines in
  let arr = Array.make_matrix size_y size_x '?' in
  List.iteri (fun j line -> String.iteri (fun i c -> arr.(j).(i) <- c) line) lines;
  arr

let rec find_index_inner elem arr pos =
  try
    if arr.(pos) = elem then
      Some pos
    else
      find_index_inner elem arr (pos + 1)
  with Invalid_argument _ -> None

let find_index elem arr =
  find_index_inner elem arr 0

let rec find_index_2d_inner elem arr pos =
  try
    match find_index elem arr.(pos) with
    | Some pos_x -> Some (pos_x, pos)
    | None -> find_index_2d_inner elem arr (pos + 1)
  with Invalid_argument _ -> None

let find_index_2d elem arr =
  find_index_2d_inner elem arr 0

let rec walk_and_mark_distance maze distances (x, y) value =
  let next_pos_list = [(x - 1, y); (x, y - 1); (x + 1, y); (x, y + 1)]
    |> List.filter (fun (x1, y1) -> maze.(y1).(x1) <> '#' && distances.(y1).(x1) < 0) in
  if List.length next_pos_list = 0 then
    [(x, y)]
  else
    let next_pos = List.hd next_pos_list in
    distances.(snd next_pos).(fst next_pos) <- value;
    (x, y) :: walk_and_mark_distance maze distances next_pos (value + 1)

let get_distance_map maze =
  let start_pos = find_index_2d 'S' maze |> Option.get in
  let distances = Array.make_matrix (Array.length maze) (Array.length maze.(0)) (-1) in
  distances.(snd start_pos).(fst start_pos) <- 0;
  let path = walk_and_mark_distance maze distances start_pos 1 in
  (distances, path)

let manhattan_distance (x1, y1) (x2, y2) =
  (Int.abs (x2 - x1)) + (Int.abs (y2 - y1))

let rate_shortcut distances (begin_x, begin_y) (end_x, end_y) = 
  try
    let begin_distance = distances.(begin_y).(begin_x) in
    let end_distance = distances.(end_y).(end_x) in
    let len = manhattan_distance (begin_x, begin_y) (end_x, end_y) in
    end_distance - begin_distance - len
  with Invalid_argument _ -> 0

let generate_offsets shortcut_length = 
  Seq.ints (-shortcut_length) |> Seq.take (shortcut_length * 2 + 1)
  |> Seq.map (fun i ->
    Seq.ints (-shortcut_length) |> Seq.take (shortcut_length * 2 + 1)
    |> Seq.map (fun j -> (i, j)))
  |> Seq.concat
  |> Seq.filter (fun (i, j) -> let dist = manhattan_distance (0, 0) (i, j) in dist > 1 && dist <= shortcut_length)
  |> List.of_seq

let get_shortcuts_from distances offsets_list (x, y) =
  offsets_list
  |> List.map (fun (i, j) -> (i + x, j + y))
  |> List.map (rate_shortcut distances (x, y))
  |> List.filter (fun x -> x > 0)

let count_cheats distances path at_least_gain max_len =
  let offsets = generate_offsets max_len in
  path
  |> List.map (get_shortcuts_from distances offsets)
  |> List.concat
  |> List.filter (fun x -> x >= at_least_gain) |> List.length 

let () =
  let f = open_in "input.txt" in
  let maze = lines_to_2d_array @@ list_of_lines f in
  let distances, path = get_distance_map maze in
  let result = count_cheats distances path 100 2 in
  let result2 = count_cheats distances path 100 20 in
  printf "%d\n%d\n" result result2