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_x size_y '?' in
  List.iteri (fun j line -> String.iteri (fun i c -> arr.(j).(i) <- c) line) lines;
  arr

let rec take_nonempty_lines lines =
  match lines with
  | head :: tail ->
    if String.length head = 0 then
      ([], tail)
    else
      let rest_of_list, remaining_lines = take_nonempty_lines tail in
      (head :: rest_of_list, remaining_lines)
  | _ -> ([], lines)

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

let add_int_pair (x1, y1) (x2, y2) =
  (x1 + x2, y1 + y2)

let rec skip_through_boxes (x, y) (dir_x, dir_y) arr =
  if arr.(y).(x) <> 'O' then
    (x, y)
  else
    let new_pos = add_int_pair (x, y) (dir_x, dir_y) in
    skip_through_boxes new_pos (dir_x, dir_y) arr

let robot_move arr pos dir =
  let dir_vector = get_dir_vector dir in
  let new_pos = add_int_pair pos dir_vector in
  let pos_behind_boxes = skip_through_boxes new_pos dir_vector arr in
  if arr.(snd pos_behind_boxes).(fst pos_behind_boxes) <> '#' then begin
    arr.(snd pos).(fst pos) <- '.';
    arr.(snd new_pos).(fst new_pos) <- '@';
    if new_pos <> pos_behind_boxes then
      arr.(snd pos_behind_boxes).(fst pos_behind_boxes) <- 'O';
    new_pos
  end else
    pos

let process_warehouse_tile c =
  match c with
  | '#' -> ['#'; '#']
  | 'O' -> ['['; ']']
  | '.' -> ['.'; '.']
  | '@' -> ['@'; '.']
  | _ -> raise (Invalid_argument "Invalid tile")

let process_warehouse arr =
  arr
    |> Array.map (fun row -> row
      |> Array.map process_warehouse_tile
      |> Array.to_seq
      |> List.of_seq
      |> List.concat
      |> List.to_seq
      |> Array.of_seq
    )

let rec try_push_horizontal (x, y) dir_x arr =
  let next_elem = arr.(y).(x) in
  if next_elem = '#' then
    false
  else if next_elem = '.' then
    true
  else
    let new_pos = (x + dir_x * 2, y) in
    if try_push_horizontal new_pos dir_x arr then begin
      let intermediate_pos = (x + dir_x, y) in
      arr.(snd new_pos).(fst new_pos) <- arr.(snd intermediate_pos).(fst intermediate_pos);
      arr.(snd intermediate_pos).(fst intermediate_pos) <- next_elem;
      true
    end else
      false

let rec possible_to_push_box_vertical (x, y) dir_y arr =
  let new_y = y + dir_y in
  possible_to_push_vertical (x, new_y) dir_y arr
  && possible_to_push_vertical (x + 1, new_y) dir_y arr

and possible_to_push_vertical (x, y) dir_y arr =
  let next_elem = arr.(y).(x) in
  if next_elem = '#' then
    false
  else if next_elem = '.' then
    true
  else if next_elem = '[' then
    possible_to_push_box_vertical (x, y) dir_y arr
  else
    possible_to_push_box_vertical (x - 1, y) dir_y arr

let rec do_push_box_vertical (x, y) dir_y arr =
  let new_y = y + dir_y in begin
  do_push_vertical (x, new_y) dir_y arr;
  do_push_vertical (x + 1, new_y) dir_y arr;
  arr.(new_y).(x) <- '[';
  arr.(new_y).(x + 1) <- ']';
  arr.(y).(x) <- '.';
  arr.(y).(x + 1) <- '.'
  end

and do_push_vertical (x, y) dir_y arr =
  let next_elem = arr.(y).(x) in
  if next_elem = '#' || next_elem = '.' then
    ()
  else if next_elem = '[' then
    do_push_box_vertical (x, y) dir_y arr
  else
    do_push_box_vertical (x - 1, y) dir_y arr

let try_push pos (dir_x, dir_y) arr =
  if dir_y = 0 then
    try_push_horizontal pos dir_x arr
  else
    if possible_to_push_vertical pos dir_y arr then begin
      do_push_vertical pos dir_y arr;
      true
    end else
      false

let robot_move_alt arr pos dir =
  let dir_vector = get_dir_vector dir in
  let new_pos = add_int_pair pos dir_vector in
  if try_push new_pos dir_vector arr then begin
    arr.(snd new_pos).(fst new_pos) <- '@';
    arr.(snd pos).(fst pos) <- '.';
    new_pos
  end else begin
    pos
  end

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 get_gps_index (x, y) =
  x + 100 * y

let () =
  let f = open_in "input.txt" in
  let layout_lines, remaining_lines = take_nonempty_lines @@ list_of_lines f in
  let arr = lines_to_2d_array layout_lines in
  let moves = String.concat "" remaining_lines in
  let pos = find_index_2d '@' arr |> Option.get in
  let warehouse1 = array_copy_2d arr in
  let _ = String.fold_left (robot_move warehouse1) pos moves in
  let result = warehouse1
    |> find_all_indices_of_2d 'O'
    |> Seq.map get_gps_index
    |> Seq.fold_left Int.add 0 in
  printf "%d\n" result;
  let warehouse2 = process_warehouse arr in
  let pos2 = find_index_2d '@' warehouse2 |> Option.get in
  let _ = String.fold_left (robot_move_alt warehouse2) pos2 moves in
  let result2 = warehouse2
    |> find_all_indices_of_2d '['
    |> Seq.map get_gps_index
    |> Seq.fold_left Int.add 0 in
  printf "%d\n" result2