defmodule Day21 do
  def run do
    lines =
      File.read!("data.txt")
      |> String.trim()
      |> String.split("\n")
      |> Enum.map(&String.graphemes/1)

    numpad = %{
      nil => {0, 0},
      "0" => {1, 0},
      "A" => {2, 0},
      "1" => {0, 1},
      "2" => {1, 1},
      "3" => {2, 1},
      "4" => {0, 2},
      "5" => {1, 2},
      "6" => {2, 2},
      "7" => {0, 3},
      "8" => {1, 3},
      "9" => {2, 3}
    }

    dpad = %{
      "<" => {0, 0},
      "v" => {1, 0},
      ">" => {2, 0},
      nil => {0, 1},
      "^" => {1, 1},
      "A" => {2, 1}
    }

    # doing some testing...
    # go([[">", "v"], ["v", ">"]], dpad) |> go(dpad) |> dbg
    # go([["<", "v"], ["v", "<"]], dpad) |> go(dpad) |> dbg
    # go([[">", "^"], ["^", ">"]], dpad) |> go(dpad) |> dbg
    # go([["<", "^"], ["^", "<"]], dpad) |> go(dpad) |> dbg
    # reveals that:
    # <, v is better than v, <
    # v, > is better than >, v
    # >, ^ is better than ^, >
    # <, ^ is better than ^, <

    part1 =
      lines
      |> go(numpad)
      |> go(dpad)
      |> go(dpad)
      |> Enum.zip(lines)
      |> Enum.map(fn {sequence, line} ->
        numeric_part =
          line |> Enum.drop(-1) |> Enum.map(&String.to_integer/1) |> Integer.undigits()

        numeric_part * length(sequence)
      end)
      |> Enum.sum()

    dbg(part1, charlists: :as_lists)

    part2 =
      lines
      |> go(numpad)
      |> Enum.zip(lines)
      # for each line
      |> Enum.reduce({%{}, 0}, fn {line, original_line}, {cache, sum} ->
        # for each key
        {new_cache, _, len} =
          Enum.reduce(line, {cache, "A", 0}, fn key, {cache, previous, sum} ->
            {num, cache} = num_steps_with_depth(previous, key, dpad, 25, cache)
            {cache, key, sum + num}
          end)

        numeric_part =
          original_line |> Enum.drop(-1) |> Enum.map(&String.to_integer/1) |> Integer.undigits()

        {new_cache, sum + numeric_part * len}
      end)

    dbg(part2)
  end

  # finds the number of steps needed to get from a to b, over and over, depth times
  defp num_steps_with_depth(a, b, keypad, depth, cache) do
    key = {a, b, depth}

    if Map.has_key?(cache, key) do
      res = Map.get(cache, key)
      {res, cache}
    else
      new_steps = steps(a, b, keypad)

      case depth do
        0 ->
          new_cache = Map.put(cache, key, 2)
          {2, new_cache}

        1 ->
          len = length(new_steps)
          new_cache = Map.put(cache, key, len)
          {len, new_cache}

        d ->
          # for each pair, find the length of depth - 1
          {new_cache, el, steps} =
            Enum.reduce(new_steps, {cache, "A", 0}, fn el, {cache, previous, num_steps} ->
              {num, new_cache} = num_steps_with_depth(previous, el, keypad, depth - 1, cache)
              new_cache = Map.put(new_cache, {previous, el, depth - 1}, num)
              {new_cache, el, num + num_steps}
            end)

          {steps, new_cache}
      end
    end
  end

  # maps over each line, then reduces over each keypress, gathering the new steps needed
  defp go(lines, pad) do
    Enum.map(lines, fn line ->
      Enum.reduce(line, {"A", []}, fn el, {previous, list} ->
        {el, Enum.concat(list, steps(previous, el, pad))}
      end)
      |> elem(1)
    end)
  end

  # returns the steps from a to b on the given keypad, including hitting A at the end
  defp steps(a, b, keypad) do
    {x1, y1} = Map.get(keypad, a)
    {x2, y2} = Map.get(keypad, b)

    y_char = if y2 > y1, do: "^", else: "v"
    x_char = if x2 > x1, do: ">", else: "<"

    xs =
      if x2 == x1,
        do: [],
        else: Enum.reduce(1..abs(x2 - x1), [], fn _el, acc -> [x_char | acc] end)

    ys =
      if y2 == y1,
        do: [],
        else: Enum.reduce(1..abs(y2 - y1), [], fn _el, acc -> [y_char | acc] end)

    {dx, dy} = Map.get(keypad, nil)

    # avoid hole
    res =
      if x2 < x1 do
        # if i'm going left, and going left _first_ hits the empty spot
        if {x2, y1} == {dx, dy} do
          # then go up/down first
          Enum.concat([ys, xs, ["A"]])
        end
      else
        # if i'm going up or down, and going up/down _first_ hits the empty spot
        if y1 != y2 and {x1, y2} == {dx, dy} do
          # then go left/right first
          Enum.concat([xs, ys, ["A"]])
        end
      end

    # if the hole needed to be avoided, return that
    if !is_nil(res) do
      res
    else
      # otherwise use our heuristic
      cond do
        # if we're going down and left, go left first
        y2 < y1 and x2 < x1 -> Enum.concat([xs, ys, ["A"]])
        # if we're going down and right, go down first
        y2 < y1 and x2 > x1 -> Enum.concat([ys, xs, ["A"]])
        # if we're going up and left, go left first
        y2 > y1 and x2 < x1 -> Enum.concat([xs, ys, ["A"]])
        # if we're going up and right (or anything else), go up first
        true -> Enum.concat([ys, xs, ["A"]])
      end
    end
  end
end

Day21.run()