(** [occurances big little] counts the number of occurances of the string
    [little] in [big]. *)
let occurances big little =
  let re = Str.regexp_string little in
  let rec impl acc pos =
    try
      let _ = Str.search_forward re big pos in
      impl (acc + 1) (Str.match_end ())
    with Not_found -> acc
  in
  impl 0 0

(** [is_inbounds sa (x, y)] returns true if (x, y) is a valid location in the
    array of strings [sa]. *)
let is_inbounds sa (x, y) =
  y >= 0 && y < Array.length sa && x >= 0 && x < String.length sa.(y)

let gen_string sa (x, y) (dx, dy) =
  let rec impl acc x y =
    if is_inbounds sa (x, y) then
      impl (String.cat acc (String.make 1 sa.(y).[x])) (x + dx) (y + dy)
    else acc
  in
  impl "" x y

let gen_strings sa (x, y) (dx, dy) (dx2, d2y) =
  let rec impl acc x y =
    if is_inbounds sa (x, y) then
      let s = gen_string sa (x, y) (dx, dy) in
      impl (s :: acc) (x + dx2) (y + d2y)
    else acc
  in
  impl [] x y

let gen_search_strings sa =
  let right = String.length sa.(0) - 1 in
  gen_strings sa (0, 0) (1, 0) (0, 1)
  @ gen_strings sa (0, 0) (0, 1) (1, 0)
  @ gen_strings sa (0, 0) (1, 1) (1, 0)
  @ gen_strings sa (0, 1) (1, 1) (0, 1)
  @ gen_strings sa (right, 0) (-1, 1) (-1, 0)
  @ gen_strings sa (right, 1) (-1, 1) (0, 1)

let find_xmas sa =
  let search_strings = gen_search_strings sa in
  List.fold_left (fun acc x -> acc + occurances x "XMAS") 0 search_strings
  + List.fold_left (fun acc x -> acc + occurances x "SAMX") 0 search_strings

let find_mas sa x y =
  let find_ms a b c d =
    sa.(y - 1).[x - 1] = a
    && sa.(y - 1).[x + 1] = b
    && sa.(y + 1).[x - 1] = c
    && sa.(y + 1).[x + 1] = d
  in
  sa.(y).[x] = 'A'
  && (find_ms 'M' 'M' 'S' 'S' || find_ms 'S' 'S' 'M' 'M'
    || find_ms 'M' 'S' 'M' 'S' || find_ms 'S' 'M' 'S' 'M')

let find_mases sa =
  let rec col_check acc x y =
    if x >= String.length sa.(y) - 1 then acc
    else if find_mas sa x y then col_check (acc + 1) (x + 1) y
    else col_check acc (x + 1) y
  in
  let rec row_check acc y =
    if y >= Array.length sa - 1 then acc
    else row_check (col_check acc 1 y) (y + 1)
  in
  row_check 0 1

let sa_of_file fname = Aoc.strings_of_file fname |> Array.of_list

let _ =
  Aoc.main sa_of_file
    [ (string_of_int, find_xmas); (string_of_int, find_mases) ]