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 file_to_2d_array in_file =
  let lines = list_of_lines in_file in
  let size_x = String.length @@ List.hd lines in
  let size_y = List.length lines in
  let arr = Array.make_matrix size_x size_y '?' in
  List.iteri (fun j line -> String.iteri (fun i c -> arr.(j).(i) <- c) line) lines;
  arr

let find_all_indices_of elem arr =
  arr |> Array.to_seq
      |> Seq.mapi (fun i x -> (i, x))
      |> Seq.filter_map (fun (i, x) -> if x = elem then Some i else None)

let find_all_indices_of_2d elem arr =
  arr |> Array.to_seq
      |> Seq.mapi (fun j row -> Seq.map (fun i -> (i, j)) (find_all_indices_of elem row))
      |> Seq.concat

let rec try_walk_to (x, y) expected_value arr visited =
  try
    if visited.(y).(x) then
      0
    else if int_of_char arr.(y).(x) <> expected_value then
      0
    else
      walk_and_count_trails (x, y) arr visited
  with Invalid_argument _ -> 0

and walk_and_count_trails (x, y) arr visited =
  visited.(y).(x) <- true;
  let elevation = int_of_char arr.(y).(x) in
  if elevation = int_of_char '9' then
    1
  else
    let next_elevation = elevation + 1 in
    try_walk_to (x - 1, y) next_elevation arr visited
    + try_walk_to (x + 1, y) next_elevation arr visited
    + try_walk_to (x, y - 1) next_elevation arr visited
    + try_walk_to (x, y + 1) next_elevation arr visited

let rec try_walk_to_no_checking (x, y) expected_value arr =
  try
    if int_of_char arr.(y).(x) <> expected_value then
      0
    else
      walk_and_count_trails_no_checking (x, y) arr
  with Invalid_argument _ -> 0

and walk_and_count_trails_no_checking (x, y) arr =
  let elevation = int_of_char arr.(y).(x) in
  if elevation = int_of_char '9' then
    1
  else
    let next_elevation = elevation + 1 in
    try_walk_to_no_checking (x - 1, y) next_elevation arr
    + try_walk_to_no_checking (x + 1, y) next_elevation arr
    + try_walk_to_no_checking (x, y - 1) next_elevation arr
    + try_walk_to_no_checking (x, y + 1) next_elevation arr

let get_trailhead_score arr pos = 
  let visited = Array.make_matrix (Array.length arr.(0)) (Array.length arr) false in
  walk_and_count_trails pos arr visited

let get_trailhead_rating arr pos = 
  walk_and_count_trails_no_checking pos arr

let () =
  let f = open_in "input.txt" in
  let arr = file_to_2d_array f in
  let trailheads = find_all_indices_of_2d '0' arr |> List.of_seq in
  let result = trailheads
    |> List.map (get_trailhead_score arr) 
    |> List.fold_left Int.add 0 in
  let result2 = trailheads
    |> List.map (get_trailhead_rating arr) 
    |> List.fold_left Int.add 0 in
  printf "%d\n" result;
  printf "%d\n" result2;