fun amountInStr (needle : string, haystack : string) : int =
  let
    val l = List.tabulate (size haystack, (fn i =>
      String.isPrefix needle (String.substring (haystack, i, (size haystack - i)))
    ))
  in
    foldl (fn (v, acc) => if v then acc + 1 else acc) 0 l
  end

fun reverseStr (str : string) = (String.implode o List.rev o String.explode) str

fun wordsInLine (line : string) =
  amountInStr ("XMAS", line) + amountInStr ("XMAS", reverseStr line)

fun wordsInLines (lines : string list) =
  foldl (op +) 0 (map wordsInLine lines)

fun columns (lines : string list) =
  let
    fun lineLength (i : int) = size (List.nth (lines, i))

    fun column (i : int) =
      String.implode (List.tabulate (
        lineLength i,
        (fn j => String.sub (List.nth (lines, j), i))
      ))
  in
    List.tabulate (length lines, column)
  end

val file = TextIO.openIn "day_4.txt"
val input = TextIO.inputAll file;
TextIO.closeIn file;

val lines = String.tokens (fn (char) => char = #"\n") input

val columns = columns lines

val result =
  wordsInLines lines +
  wordsInLines columns