open Printf;;
open String;;

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 = 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 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 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 find_guard_position arr =
  arr |> find_index_2d '^' |> Option.get

type dir =
  | Up
  | Down
  | Left
  | Right

let get_new_pos guard_pos dir =
  match dir with
  | Up -> (fst guard_pos, snd guard_pos - 1)
  | Down -> (fst guard_pos, snd guard_pos + 1)
  | Left -> (fst guard_pos - 1, snd guard_pos)
  | Right -> (fst guard_pos + 1, snd guard_pos)

let turn_right dir =
  match dir with
  | Up -> Right
  | Down -> Left
  | Left -> Up
  | Right -> Down

let int_of_dir dir =
  match dir with
  | Up -> 0
  | Down -> 1
  | Left -> 2
  | Right -> 3

let guard_state_index pos dir =
  Int.logor (fst pos) (
  Int.logor (Int.shift_left (snd pos) 12)
            (Int.shift_left (int_of_dir dir) 24))

module IntSet = Set.Make(Int)

let rec rotate_until_can_move arr guard_pos dir =
  let new_pos = get_new_pos guard_pos dir in
  if arr.(snd new_pos).(fst new_pos) <> '#' then
    dir
  else
    let new_dir = turn_right dir in
    rotate_until_can_move arr guard_pos new_dir

(* returns true if guard is able to complete the walk, false if guard falls into a cycle *)
let rec guard_walk_inner arr previous_states guard_pos dir =
  try
    arr.(snd guard_pos).(fst guard_pos) <- 'X';
    let guard_state_index = (guard_state_index guard_pos dir) in
    if Option.is_some @@ IntSet.find_opt guard_state_index previous_states then
      false
    else
      let new_previous_states = IntSet.add guard_state_index previous_states in
      let new_dir = rotate_until_can_move arr guard_pos dir in
      let new_pos = get_new_pos guard_pos new_dir in
      guard_walk_inner arr new_previous_states new_pos new_dir
  with Invalid_argument _ -> true

let guard_walk arr guard_pos dir =
  guard_walk_inner arr IntSet.empty guard_pos dir

let count_2d value arr =
  Array.fold_left 
    (fun acc elem -> acc + 
      Array.fold_left 
        (fun acc2 elem2 -> acc2 + if elem2 = value then 1 else 0)
        0 elem) 
    0 arr

let array_copy_2d arr =
  Array.init (Array.length arr) (fun i -> Array.copy arr.(i))

let try_block_guard arr guard_pos obstacle_pos =
  let array_copy = array_copy_2d arr in
  array_copy.(snd obstacle_pos).(fst obstacle_pos) <- '#';
  not @@ guard_walk array_copy guard_pos Up

let () =
  let f = open_in "input.txt" in
  let arr = file_to_2d_array f in
  let guard_pos = find_guard_position arr in
  let _ = guard_walk arr guard_pos Up in
  let visited = List.of_seq (find_all_indices_of_2d 'X' arr) in
  let result = List.length visited in
  printf "%d\n%!" result;
  let cycle_positions = List.filter (try_block_guard arr guard_pos) visited in
  let result2 = List.length cycle_positions in
  printf "%d\n" result2;