defmodule Day04 do
  def part1(data, word \\ ~c"XMAS") do
    data = data
    |> String.split("\n", trim: true)
    |> Enum.map(&(:array.from_list(String.to_charlist(&1))))

    count_word(0, data, word) +
      count_word(0, data, Enum.reverse(word))
  end

  def count_word(acc, [], _) do
    acc
  end

  def count_word(acc, data, word) do
    count_word(acc + count_word_hd(data, word), tl(data), word)
  end

  def count_word_hd(data, word) do
    0..:array.size(hd(data))-1
    |> Enum.map(&(count_word_hd_col(data, word, &1)))
    |> Enum.sum()
  end

  def count_word_hd_col(data, word, col) do
    [check_word_hori(data, word, col),
     check_word_vert(data, word, col),
     check_word_diag_right(data, word, col),
     check_word_diag_left(data, word, col)]
    |> Enum.count(&(&1))
  end

  def check_word_hori(data, word, col) do
    check_word(data, word, col, &(&1), &(&1 + 1))
  end

  def check_word_vert(data, word, col) do
    check_word(data, word, col, &tl/1, &(&1))
  end

  def check_word_diag_right(data, word, col) do
    check_word(data, word, col, &tl/1, &(&1 + 1))
  end

  def check_word_diag_left(data, word, col) do
    check_word(data, word, col, &tl/1, &(&1 - 1))
  end

  def check_word(data, word, col, data_step_fn, col_step_fn) do
    cond do
      word == [] ->
        true
      data == [] or col < 0 or col >= :array.size(hd(data)) ->
        false
      hd(word) == :array.get(col, hd(data)) ->
        check_word(
          data_step_fn.(data),
          tl(word),
          col_step_fn.(col),
          data_step_fn,
          col_step_fn
        )
      true ->
        false
    end
  end

  def part2(data, word \\ ~c"MAS") do
    data = data
    |> String.split("\n", trim: true)
    |> Enum.map(&(:array.from_list(String.to_charlist(&1))))

    count_cross(0, data, word)
  end

  def count_cross(acc, [], _) do
    acc
  end

  def count_cross(acc, data, word) do
    count_cross(acc + count_cross_hd(data, word), tl(data), word)
  end

  def count_cross_hd(data, word) do
    0..:array.size(hd(data))-1
    |> Enum.count(&(check_cross_hd_col(data, word, &1)))
  end

  def check_cross_hd_col(data, word, col) do
    (check_word_diag_right(data, word, col) and
       check_word_diag_left(data, word, col+2)) or
      (check_word_diag_right(data, Enum.reverse(word), col) and
         check_word_diag_left(data, word, col+2)) or
      (check_word_diag_right(data, word, col) and
         check_word_diag_left(data, Enum.reverse(word), col+2)) or
      (check_word_diag_right(data, Enum.reverse(word), col) and
         check_word_diag_left(data, Enum.reverse(word), col+2))
  end
end

data = IO.read(:stdio, :eof)

{time1 , ans1} = :timer.tc(fn -> Day04.part1(data) end)
IO.puts("Time  : #{time1 / 1000000}")
IO.puts("Answer: #{ans1}")

{time2 , ans2} = :timer.tc(fn -> Day04.part2(data) end)
IO.puts("Time  : #{time2 / 1000000}")
IO.puts("Answer: #{ans2}")