defmodule Aoc2023.Day17 do
  use Aoc2023

  defmodule PriorityQueue do
    defstruct [:set]
    def new(contents \\ []), do: %__MODULE__{set: :gb_sets.from_list(contents)}

    def add(%__MODULE__{set: set}, elem, prio) do
      %__MODULE__{set: :gb_sets.add({prio, elem}, set)}
    end

    def size(%__MODULE__{set: set}) do
      :gb_sets.size(set)
    end

    def pop_min(%__MODULE__{set: set}) do
      case :gb_sets.size(set) do
        0 ->
          {nil, new()}

        _ ->
          {res, set} = :gb_sets.take_smallest(set)
          {res, %__MODULE__{set: set}}
      end
    end

    defimpl Inspect do
      import Inspect.Algebra

      def inspect(%PriorityQueue{set: set}, opts) do
        concat(["#PriorityQueue.new(", to_doc(:gb_sets.to_list(set), opts), ")"])
      end
    end
  end

  def parse(input) do
    lines = input |> lines

    {
      {lines |> length() |> Kernel.-(1), lines |> head |> String.length() |> Kernel.-(1)},
      lines
      |> Enum.with_index()
      |> Enum.flat_map(fn {row, y} ->
        row
        |> String.to_charlist()
        |> Enum.with_index()
        |> Enum.map(fn {cost, x} -> {{y, x}, String.to_integer(<<cost>>)} end)
      end)
      |> Enum.into(%{})
    }
  end

  defp a_star(starts, neighbors, heuristic, is_goal) do
    loop = fn recur, explored, frontier ->
      {explored, frontier}
      {{_, current}, frontier} = PriorityQueue.pop_min(frontier)

      if is_goal.(current) do
        {explored, current}
      else
        cost_to_here = Map.get(explored, current) |> elem(1)

        {explored, frontier} =
          for {neighbor, cost} <- neighbors.(current), reduce: {explored, frontier} do
            {explored, frontier} ->
              cost_ = cost_to_here + cost

              if Map.get(explored, neighbor, {nil, cost_}) |> elem(1) < cost_ do
                {explored, frontier}
              else
                {
                  Map.put(explored, neighbor, {current, cost_}),
                  PriorityQueue.add(frontier, neighbor, cost_ + heuristic.(neighbor))
                }
              end
          end

        recur.(recur, explored, frontier)
      end
    end

    {explored, goal} =
      loop.(
        loop,
        starts |> Enum.map(&{&1, {nil, 0}}) |> Enum.into(%{}),
        starts |> Enum.map(&{0, &1}) |> PriorityQueue.new()
      )

    # This doesn't include the starting point
    path_find = fn recur, point, path ->
      case Map.get(explored, point) do
        nil -> path
        {nil, _} -> path
        {prev, _} -> recur.(recur, prev, [point | path])
      end
    end

    path_find.(path_find, goal, [])
  end

  def part1({{my, mx}, grid}) do
    # Nodes are y, x, dir, straight length
    a_star(
      [{0, 0, :d, 0}, {0, 0, :r, 0}],
      fn {y, x, dir, len} ->
        case dir do
          :u ->
            [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}] ++
              if len < 2, do: [{y + 1, x, dir, len + 1}], else: []

          :d ->
            [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}] ++
              if len < 2, do: [{y - 1, x, dir, len + 1}], else: []

          :l ->
            [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}] ++
              if len < 2, do: [{y, x - 1, dir, len + 1}], else: []

          :r ->
            [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}] ++
              if len < 2, do: [{y, x + 1, dir, len + 1}], else: []
        end
        |> Enum.filter(fn {y, x, _, _} -> y in 0..my and x in 0..mx end)
        |> Enum.map(fn {y, x, _, _} = pos -> {pos, Map.get(grid, {y, x})} end)
      end,
      fn {y, x, _, _} -> abs(y - my) + abs(x - mx) end,
      fn {y, x, _, _} -> {y, x} == {my, mx} end
    )
    |> Enum.map(fn {y, x, _, _} -> Map.get(grid, {y, x}) end)
    |> Enum.sum()
  end

  def part2({{my, mx}, grid}) do
    # Nodes are y, x, dir, straight length
    a_star(
      [{0, 0, :r, 0}, {0, 0, :d, 0}, {0, 0, :u, 0}, {0, 0, :l, 0}],
      fn {y, x, dir, len} ->
        (if(len < 9,
           do:
             case dir do
               :u -> [{y + 1, x, :u, len + 1}]
               :d -> [{y - 1, x, :d, len + 1}]
               :l -> [{y, x - 1, :l, len + 1}]
               :r -> [{y, x + 1, :r, len + 1}]
             end,
           else: []
         ) ++
           if(len > 2,
             do:
               case dir do
                 :u -> [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}]
                 :d -> [{y, x - 1, :l, 0}, {y, x + 1, :r, 0}]
                 :l -> [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}]
                 :r -> [{y + 1, x, :u, 0}, {y - 1, x, :d, 0}]
               end,
             else: []
           ))
        |> Enum.filter(fn {y, x, _, _} -> y in 0..my and x in 0..mx end)
        |> Enum.map(fn {y, x, _, _} = pos -> {pos, Map.get(grid, {y, x})} end)
      end,
      fn {y, x, _, _} -> abs(y - my) + abs(x - mx) end,
      fn {y, x, _, len} -> {y, x} == {my, mx} and len > 2 end
    )
    |> Enum.map(fn {y, x, _, _} -> Map.get(grid, {y, x}) end)
    |> Enum.sum()
  end
end