fun amountInStr (needle, haystack) : 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.implode o List.rev o String.explode) str

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

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

fun lineLength lines i = size (List.nth (lines, i))

fun pos lines (x, y) = String.sub (List.nth (lines, y), x)

fun column lines i =
  String.implode (List.tabulate (
    lineLength lines i,
    (fn j => pos lines (i, j))
  ))

fun nwDiagonal lines i =
  String.implode (List.tabulate (i + 1, (fn j =>
    pos lines (i - j, j)
  )))

fun seDiagonal lines i =
  String.implode (List.tabulate (i, (fn j =>
    pos lines ((lineLength lines i) - (i - j), (lineLength lines i) - j - 1)
  )))

fun neDiagonal lines i =
  String.implode (List.tabulate (i, (fn j =>
    pos lines ((lineLength lines i) - (i - j), j)
  )))

fun swDiagonal lines i =
  String.implode (List.tabulate (i + 1, (fn j =>
    pos lines (i - j, (lineLength lines i) - j - 1)
  )))

fun applyToLines f lines =
  List.tabulate (length lines, (f lines))

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

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

val columns = applyToLines column lines

val nwDiagonals = applyToLines nwDiagonal lines
val seDiagonals = applyToLines seDiagonal lines
val neDiagonals = applyToLines neDiagonal lines
val swDiagonals = applyToLines swDiagonal lines

val result =
  wordsInLines lines +
  wordsInLines columns +
  wordsInLines nwDiagonals +
  wordsInLines seDiagonals +
  wordsInLines neDiagonals +
  wordsInLines swDiagonals