let string_to_char_list s = 
  s |> String.to_seq |> List.of_seq
let lines = Reuse.split "inputs/day04.txt" '\n'
let input = List.map string_to_char_list lines

let (>>=) = Option.bind ;;
let get_grid_pos (x: int) (y: int) (grid: char list list) : char option =
  List.nth_opt grid x >>= fun sublist -> List.nth_opt sublist y

let get_neighbors (x: int) (y: int) (grid: char list list) : char option list =
  let right_top = if y = 0 then None else get_grid_pos (x + 1) (y - 1) grid in
  let left = if x = 0 then None else get_grid_pos (x - 1) y grid in
  let left_bot = if x = 0 then None else get_grid_pos (x - 1) (y + 1) grid in
  let top = if y = 0 then None else get_grid_pos x (y - 1) grid in
  let left_top = if y = 0 || x = 0 then None else get_grid_pos (x - 1) (y - 1) grid in
  [get_grid_pos (x + 1) y grid; get_grid_pos (x + 1) (y + 1) grid; right_top; left; left_bot; left_top;
  get_grid_pos x (y + 1) grid; top]

let sum_papers (neighbors: char option list) : int =
  List.length (List.filter (fun x -> x = Some '@') neighbors)

let update_grid_position (grid: char list list) (x, y: int * int) : char list list =
  let update_row (row: char list) (y: int) : char list =
    List.mapi (fun col_idx c -> if col_idx = y then '.' else c) row
  in
  List.mapi (fun row_idx row ->
    if row_idx = x then update_row row y else row
  ) grid

let apply_to_all_positions (grid: char list list) (positions: (int * int) list) : char list list =
  List.fold_left (fun current_grid (x,y) -> update_grid_position current_grid (x, y)) grid positions

let rec update_grid (grid: char list list) (acc : int) : int =
  let result = ref acc in
  let rows = List.length grid in
  let cols = List.length (List.hd grid) in
  let positions = ref [] in
  for x = 0 to rows - 1 do
    for y = 0 to cols - 1 do
    let ats = sum_papers (get_neighbors x y grid) in
    if get_grid_pos x y grid = Some '@' && ats < 4 then positions := [(x, y)] @ !positions;
    result := !result + List.length !positions;
  done;
  done;
  let new_grid = apply_to_all_positions grid !positions in
  if new_grid = grid then acc else update_grid new_grid (acc + List.length !positions)

let solve_part1 (inputs: char list list) : int =
  let result = ref 0 in
  let rows = List.length inputs in
  let cols = List.length (List.hd inputs) in
  for x = 0 to rows - 1 do
    for y = 0 to cols - 1 do
    let ats = sum_papers (get_neighbors x y inputs) in
    if get_grid_pos x y inputs = Some '@' && ats < 4 then result := !result + 1
  done;
  done;
  !result


let solve_part2 (inputs: char list list) : int =
  update_grid inputs 0

let day4_part1 : int =
  solve_part1 input

let day4_part2 : int = 
  solve_part2 input

module Tests = struct
  let test_input = "..@@.@@@@.
@@@.@.@.@@
@@@@@.@.@@
@.@@@@..@.
@@.@@@@.@@
.@@@@@@@.@
.@.@.@.@@@
@.@@@.@@@@
.@@@@@@@@.
@.@.@@@.@."
  
  
  
  let run_tests () = 
    let test_lines = String.split_on_char '\n' test_input in
    let test_input = List.map string_to_char_list test_lines in

    let part1_result = solve_part1 test_input in
    let part2_result = solve_part2 test_input in
    if part1_result <> 13 then
      failwith ("Day 4 Part 1 test failed: expected 13, got " ^ string_of_int part1_result);
      
   if part2_result <> 43 then
        failwith ("Day 4 Part 2 test failed: expected 43, got " ^ string_of_int part2_result);
      
    Printf.printf "All day 4 tests passed!\n"
end
  
let () = Tests.run_tests()