defmodule Day21 do
  defmodule Keypad do
    @derive Inspect
    @enforce_keys [:type]
    defstruct @enforce_keys

    def new(type) do
      %Keypad{type: type}
    end

    def shortest_paths(keypad, from, to) do
      bfs([[{:nil, from}]], to, keypad)
    end

    defp bfs(paths, to, keypad) do
      finished = Enum.any?(paths, &(elem(hd(&1), 1) == to))

      if finished do
        Enum.filter(paths, &(elem(hd(&1), 1) == to))
      else
        for path <- paths, neighbor <- neighbors(keypad, elem(hd(path), 1)) do
          [neighbor | path]
        end
        |> bfs(to, keypad)
      end
    end

    defp neighbors(keypad, from) do
      case keypad.type do
        :numeric -> neighbors_numeric(from)
        :directional -> neighbors_directional(from)
      end
    end

    defp neighbors_numeric(from) do
      case from do
        "A" -> %{left: "0", up: "3"}
        "0" -> %{right: "A", up: "2"}
        "1" -> %{right: "2", up: "4"}
        "2" -> %{down: "0", left: "1", right: "3", up: "5"}
        "3" -> %{down: "A", left: "2", up: "6"}
        "4" -> %{down: "1", right: "5", up: "7"}
        "5" -> %{down: "2", left: "4", right: "6", up: "8"}
        "6" -> %{down: "3", left: "5", up: "9"}
        "7" -> %{down: "4", right: "8"}
        "8" -> %{down: "5", left: "7", right: "9"}
        "9" -> %{down: "6", left: "8"}
      end
    end

    defp neighbors_directional(from) do
      case from do
        :A -> %{left: :up, down: :right}
        :up -> %{right: :A, down: :down}
        :left -> %{right: :down}
        :down -> %{left: :left, up: :up, right: :right}
        :right -> %{left: :down, up: :A}
      end
    end
  end

  def part1(data, n \\ 2) do
    data
    |> String.split("\n", trim: true)
    |> Enum.map(fn code ->
      {
        String.split(code, "A") |> hd() |> String.to_integer(),
        shortest_dist(String.codepoints(code), "A", n, %{}) |> elem(0)
      }
    end)
    # |> IO.inspect(charlists: :as_lists)
    |> Enum.map(fn {a, b} -> a*b end)
    |> Enum.sum()
  end

  defp shortest_dist(code, start, n, state) do
    chunks = [start | code]
    |> Enum.chunk_every(2, 1, :discard)

    kpadtype = if start == "A" do :numeric else :directional end

    {values, nstate} =
      for [from, to] <- chunks, reduce: {[], state} do
        {values, state} ->
          shortest_dist(from, to, n, kpadtype, state)
          |> then(fn {v, state} -> {[v | values], state} end)
      end

    {Enum.sum(values), nstate}
  end

  defp shortest_dist(from, to, n, keypadtype, state) do
    key = {from, to, n, keypadtype}
    case Map.get(state, key) do
      :nil ->
        {values, nstate} =
          for path <- paths(from, to, keypadtype), reduce: {[], state} do
            {values, state} -> shortest_dist(path, n, state)
              |> then(fn {v, s} -> {[v | values], s} end)
          end
        value = Enum.min(values)
        {value, Map.put(nstate, key, value)}
      value ->
        {value, state}
    end
  end

  defp paths(from, to, keypad) do
    case keypad do
      :numeric ->
        Keypad.new(:numeric)
      :directional -> 
        Keypad.new(:directional)
    end
    |> Keypad.shortest_paths(from, to)
  end

  defp shortest_dist(path, n, state) do
    if n == 0 do
      {Enum.count(path), state}
    else
      path
      |> Enum.map(&(elem(&1, 0)))
      |> then(&([:A | &1]))
      |> Enum.reverse()
      |> tl()
      |> shortest_dist(:A, n-1, state)
    end
  end

  def part2(data) do
    part1(data, 25)
  end
end

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

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

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