defmodule Day20 do
  defmodule Maze do
    @derive Inspect
    @enforce_keys [:positions, :walls, :source, :target]
    defstruct @enforce_keys

    def from_data(data) do
      positions = data
      |> String.split("\n", trim: true)
      |> Enum.with_index()
      |> Enum.map(fn {l, i} ->
        l
        |> String.codepoints()
        |> Enum.with_index()
        |> Enum.map(fn {v, j} -> {{i, j}, v} end)
      end)
      |> List.flatten()

      source = Enum.find(positions, fn {_, v} -> v == "S" end) |> elem(0)
      target = Enum.find(positions, fn {_, v} -> v == "E" end) |> elem(0)

      walls = positions
      |> Enum.filter(fn {_, v} -> v == "#" end)
      |> Enum.map(fn {p, _} -> p end)

      positions = positions
      |> Enum.filter(fn {_, v} -> v != "#" end)
      |> Enum.map(fn {p, _} -> p end)
      |> MapSet.new()

      %Maze{positions: positions, walls: walls, source: source, target: target}
    end

    def cheats(maze, seconds) do
      # Get shortest distances for all nodes
      dist = rev_shortest(maze)

      # Get cheats
      for n1 <- maze.positions, n2 <- neighbors(maze, n1, seconds) do
        {n1, n2, Map.get(dist, n1) - Map.get(dist, n2) - dist(n1, n2)}
      end
    end

    defp dist({x1, y1}, {x2, y2}), do: abs(x1-x2) + abs(y1-y2)

    defp rev_shortest(maze) do
      maze.target
      |> then(&rev_shortest(maze, :queue.from_list([{&1, 0}]), MapSet.new([&1]), Map.new()))
    end

    defp rev_shortest(maze, q, inq, dist) do
      case :queue.peek(q) do
        :empty ->
          dist
        {_, {v, d}} ->
          for n <- neighbors(maze, v), reduce: {:queue.drop(q), inq} do
            {q, inq} ->
              if MapSet.member?(inq, n) do
                {q, inq}
              else
                {:queue.in({n, d+1}, q), MapSet.put(inq, n)}
              end
          end
          |> then(fn {q, inq} -> rev_shortest(maze, q, inq, Map.put(dist, v, d)) end)
      end
    end

    defp neighbors(maze, {x, y}) do
      [{x+1, y},
       {x-1, y},
       {x, y+1},
       {x, y-1}]
      |> Enum.filter(&MapSet.member?(maze.positions, &1))
    end

    defp neighbors(maze, {x, y}, d) do
      for i <- 0..d, j <- 0..d-i do
        cond do
          i == 0 and j == 0 ->
            []
          i == 0 ->
            [{x, y+j},
             {x, y-j}]
          j == 0 ->
            [{x+i, y},
             {x-i, y}]
          true ->
            [{x+i, y+j},
             {x-i, y+j},
             {x+i, y-j},
             {x-i, y-j}]
        end
      end
      |> List.flatten()
      |> Enum.filter(&MapSet.member?(maze.positions, &1))
    end
  end

  def part1(data) do
    data
    |> Maze.from_data()
    |> Maze.cheats(2)
    |> Enum.count(& elem(&1, 2) >= 100)
  end

  def part2(data) do
    data
    |> Maze.from_data()
    |> Maze.cheats(20)
    |> Enum.count(& elem(&1, 2) >= 100)
  end
end

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

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

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