open Printf;;
open Buffer;;
open String;;
open Str;;

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 count_in_string_inner re str pos acc =
  try
    let new_pos = Str.search_forward re str pos in
    count_in_string_inner re str (new_pos + 1) (acc + 1)
  with Not_found -> acc

let count_in_string str =
  let re = Str.regexp "XMAS\\|SAMX" in
  count_in_string_inner re str 0 0

let prepare_horizontal_str arr =
  let buf = Buffer.create 1024 in
  Array.iter (fun row -> Array.iter (fun c -> Buffer.add_char buf c) row; Buffer.add_char buf '\n') arr;
  Buffer.contents buf

let prepare_vertical_str arr =
  let buf = Buffer.create 1024 in
  Seq.iter (fun i -> Array.iter (fun row -> Buffer.add_char buf row.(i)) arr; Buffer.add_char buf '\n')
    (Seq.init (Array.length arr.(0)) (fun x -> x));
  Buffer.contents buf

let rec traverse_diagonal fn arr x y  =
  try
    let () = fn arr.(y).(x) in
    traverse_diagonal fn arr (x + 1) (y + 1)
  with Invalid_argument _ -> ()

let rec traverse_antidiagonal fn arr x y  =
  try
    let () = fn arr.(y).(x) in
    traverse_antidiagonal fn arr (x - 1) (y + 1)
  with Invalid_argument _ -> ()

let range start stop =
  Seq.init (stop - start) (fun x -> start + x)

let range_rev start stop =
  Seq.init (stop - start) (fun x -> stop - 1 - x)

let prepare_diagonal_str arr =
  let buf = Buffer.create 1024 in
  let length_x = Array.length arr.(0) in
  let length_y = Array.length arr in
  Seq.iter (fun y -> traverse_diagonal (Buffer.add_char buf) arr 0 y; Buffer.add_char buf '\n')
    (range_rev 1 (length_y - 1));
  Seq.iter (fun x -> traverse_diagonal (Buffer.add_char buf) arr x 0; Buffer.add_char buf '\n')
    (range 0 length_x);
  Buffer.contents buf

let prepare_antidiagonal_str arr =
  let buf = Buffer.create 1024 in
  let length_x = Array.length arr.(0) in
  let length_y = Array.length arr in
  Seq.iter (fun x -> traverse_antidiagonal (Buffer.add_char buf) arr x 0; Buffer.add_char buf '\n')
    (range 0 length_x);
  Seq.iter (fun y -> traverse_antidiagonal (Buffer.add_char buf) arr (length_x - 1) y; Buffer.add_char buf '\n')
    (range 1 (length_y - 1));
  Buffer.contents buf

let is_mas_x arr x y =
  if arr.(y).(x) = 'A' then
    let nw = arr.(y - 1).(x - 1) in
    let se = arr.(y + 1).(x + 1) in
    if (nw = 'M' && se = 'S') || (nw = 'S' && se = 'M') then
      let ne = arr.(y - 1).(x + 1) in
      let sw = arr.(y + 1).(x - 1) in
      (ne = 'M' && sw = 'S') || (ne = 'S' && sw = 'M')
    else
      false
  else
    false

let count_mas_x arr =
  let length_x = Array.length arr.(0) in
  let length_y = Array.length arr in
  Seq.fold_left 
    (fun acc y -> acc +
      Seq.fold_left 
        (fun acc x -> acc +
          if is_mas_x arr x y then 1 else 0)
        0 (range 1 (length_x - 1)))
    0 (range 1 (length_y - 1))

let () =
  let f = open_in "input.txt" in
  let arr = file_to_2d_array f in
  let counter = 
    (count_in_string @@ prepare_horizontal_str arr)
    + (count_in_string @@ prepare_vertical_str arr)
    + (count_in_string @@ prepare_diagonal_str arr)
    + (count_in_string @@ prepare_antidiagonal_str arr) in
  let counter2 = count_mas_x arr in
  printf "%d\n" counter;
  printf "%d\n" counter2