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 print_layout arr =
  arr
    |> Array.iter (fun row -> row
      |> Array.iter (fun x -> print_char x);
      print_endline "")

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 add_int_pair (x1, y1) (x2, y2) =
  (x1 + x2, y1 + y2)

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

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 chain_comparison comp1 comp2 arg1 arg2 =
  match comp1 arg1 arg2 with
  | 0 -> comp2 arg1 arg2
  | c -> c

type maze_walk_state = {
  pos: (int * int);
  dir: char;
  score: int
}

module MazeWalkState =
struct
  type t = maze_walk_state
  let compare arg0 arg1 =
    chain_comparison (fun arg1 arg2 -> Stdlib.compare arg1.score arg2.score) (
      chain_comparison (fun arg1 arg2 -> Stdlib.compare (fst arg1.pos) (fst arg2.pos)) (
        chain_comparison (fun arg1 arg2 -> Stdlib.compare (snd arg1.pos) (snd arg2.pos))
                           (fun arg1 arg2 -> Stdlib.compare arg1.dir arg2.dir)
      )
    ) arg0 arg1
end
module MazeWalkStateSet = Set.Make(MazeWalkState)
module IntMap = Map.Make(Int)
module IntSet = Set.Make(Int)

let get_dir_vector dir =
  match dir with
  | '<' -> (-1, 0)
  | '>' -> (1, 0)
  | '^' -> (0, -1)
  | 'v' -> (0, 1)
  | _ -> raise (Invalid_argument "Invalid direction")

let rotate_dir_right dir =
  match dir with
  | '<' -> '^'
  | '>' -> 'v'
  | '^' -> '>'
  | 'v' -> '<'
  | _ -> raise (Invalid_argument "Invalid direction")

let rotate_dir_left dir =
  match dir with
  | '<' -> 'v'
  | '>' -> '^'
  | '^' -> '<'
  | 'v' -> '>'
  | _ -> raise (Invalid_argument "Invalid direction")

let walk_forward state =
  {state with pos = add_int_pair state.pos (get_dir_vector state.dir); score = state.score + 1 }

let turn_right state =
  walk_forward {state with score = state.score + 1000; dir = rotate_dir_right state.dir}
 
let turn_left state =
  walk_forward {state with score = state.score + 1000; dir = rotate_dir_left state.dir}

let dir_index dir =
  match dir with
  | '<' -> 0
  | '>' -> 1
  | '^' -> 2
  | 'v' -> 3
  | _ -> raise (Invalid_argument "Invalid direction")

let generate_next_moves maze state =
  [walk_forward state; turn_left state; turn_right state]
    |> List.filter (fun state -> let (x, y) = state.pos in maze.(y).(x) <> '#')

let get_state_index maze {pos=(x, y); dir=dir; score=_} =
  ((Array.length maze) * x + y) * 4 + (dir_index dir)

let state_index_remove_rotation state_index =
  state_index / 4

let decode_position_index maze position_index = 
  let height = Array.length maze in
  (position_index / height, position_index mod height)

let add_prev maze prev_state next_state prev_move_map =
  let prev_index = get_state_index maze prev_state in
  let next_index = get_state_index maze next_state in
  let update_fun lst_opt =
    match lst_opt with
    | Some (score, lst) ->
        if next_state.score = score then
          Some (score, prev_index :: lst)
        else if next_state.score < score then
          Some (next_state.score, [prev_index])
        else
          Some (score, lst)
    | None -> Some (next_state.score, [prev_index]) in
  IntMap.update next_index update_fun prev_move_map

let worse_than_final_score final_score_opt score =
  match final_score_opt with
  | Some final_score -> score > final_score
  | None -> false

let rec dijkstra_find_path maze visited_array end_pos final_score_opt prev_move_map next_move_queue =
  match MazeWalkStateSet.min_elt_opt next_move_queue with
  | Some state ->
    let queue_with_removed = MazeWalkStateSet.remove state next_move_queue in
    if visited_array.(get_state_index maze state) then
      dijkstra_find_path maze visited_array end_pos final_score_opt prev_move_map queue_with_removed
    else
      if worse_than_final_score final_score_opt state.score then
        (Option.get final_score_opt, prev_move_map)
      else begin
        visited_array.(get_state_index maze state) <- true;
        let next_states = generate_next_moves maze state in
        let new_prev_move_map = List.fold_left (fun map s -> add_prev maze state s map) prev_move_map next_states in
        let solution_states = List.filter (fun s -> s.pos = end_pos) next_states in
        if List.length solution_states > 0 then
          let final_score = (List.hd solution_states).score in
          dijkstra_find_path maze visited_array end_pos (Some final_score) new_prev_move_map queue_with_removed
        else
          let new_queue = List.fold_left (fun queue s -> MazeWalkStateSet.add s queue) queue_with_removed next_states in
          dijkstra_find_path maze visited_array end_pos final_score_opt new_prev_move_map new_queue
    end
  | None -> (-1, prev_move_map)

let rec backtrack_and_fill_position_set prev_move_map state_index pos_set =
  let new_pos_set = IntSet.add (state_index_remove_rotation state_index) pos_set in
  match IntMap.find_opt state_index prev_move_map with
  | Some (_, prev_list) ->
    List.fold_left (fun set idx -> backtrack_and_fill_position_set prev_move_map idx set) new_pos_set prev_list
  | None -> new_pos_set

let count_tiles_on_best_paths maze end_pos prev_move_map =
  ['<'; '>'; '^'; 'v']
  |> List.map (fun dir -> get_state_index maze {pos=end_pos; dir=dir; score=0})
  |> List.fold_left (fun set state_index -> backtrack_and_fill_position_set prev_move_map state_index set) IntSet.empty

let find_best_path maze start_pos end_pos =
  let visited_array = Array.make ((Array.length maze) * (Array.length maze.(0) * 4)) false in
  let start_state = {pos=start_pos; dir='>'; score=0} in
  dijkstra_find_path maze visited_array end_pos None IntMap.empty (MazeWalkStateSet.add start_state MazeWalkStateSet.empty)

let () =
  let f = open_in "input.txt" in
  let maze = lines_to_2d_array @@ list_of_lines f in
  let start_pos = find_index_2d 'S' maze |> Option.get in
  let end_pos = find_index_2d 'E' maze |> Option.get in
  let (result, prev_move_map) = find_best_path maze start_pos end_pos in
  let result2 = prev_move_map |> count_tiles_on_best_paths maze end_pos |> IntSet.cardinal in
  printf "%d\n%d\n" result result2