defmodule AOC.Day17 do
  use AOC.Day, day: 17
  # use AOC.Day, day: 17, input: "example2"

  def parse_input(lines) do
    lines
    |> Enum.with_index()
    |> Enum.flat_map(fn {line, y} ->
      line
      |> String.split("", trim: true)
      |> Enum.with_index()
      |> Enum.map(fn {heat_loss, x} ->
        {{y, x}, String.to_integer(heat_loss)}
      end)
    end)
    |> Enum.into(%{})
  end

  def turn(:north, :left), do: :west
  def turn(:north, :right), do: :east
  def turn(:east, :left), do: :north
  def turn(:east, :right), do: :south
  def turn(:south, :left), do: :east
  def turn(:south, :right), do: :west
  def turn(:west, :left), do: :south
  def turn(:west, :right), do: :north

  def apply_direction({y, x}, :north), do: {y - 1, x}
  def apply_direction({y, x}, :east), do: {y, x + 1}
  def apply_direction({y, x}, :south), do: {y + 1, x}
  def apply_direction({y, x}, :west), do: {y, x - 1}

  def simulate_movement(map, goal, min_straight, max_straight) do
    simulate_movement(map, goal, min_straight, max_straight, {0, 1}, :east, 1, 0, %{})
    simulate_movement(map, goal, min_straight, max_straight, {1, 0}, :south, 1, 0, %{})
  end

  def simulate_movement(
        _map,
        _goal,
        _,
        _,
        position,
        _direction,
        _straight_count,
        _heat_loss,
        visited
      )
      when is_map_key(visited, position),
      do: :ok

  def simulate_movement(
        map,
        goal,
        min_straight,
        max_straight,
        position,
        direction,
        straight_count,
        heat_loss,
        visited
      ) do
    case Map.fetch(map, position) do
      {:ok, extra_heat_loss} ->
        heat_loss = heat_loss + extra_heat_loss

        better_heat_loss? =
          case :ets.lookup(:movement_memo, {position, direction, straight_count}) do
            [] ->
              true

            [{_, heat_loss_memo}] ->
              heat_loss < heat_loss_memo
          end

        if better_heat_loss? do
          if position == goal do
            if straight_count >= min_straight do
              :ets.insert(:movement_memo, {{position, direction, straight_count}, heat_loss})

              case :ets.lookup(:best_memo, :best) do
                [] ->
                  IO.puts("BEST: #{heat_loss}")
                  :ets.insert(:best_memo, {:best, heat_loss})

                [{:best, best}] when heat_loss < best ->
                  IO.puts("BEST: #{heat_loss}")
                  :ets.insert(:best_memo, {:best, heat_loss})

                _ ->
                  :ok
              end
            end
          else
            :ets.insert(:movement_memo, {{position, direction, straight_count}, heat_loss})

            worse_than_best =
              case :ets.lookup(:best_memo, :best) do
                [] -> false
                [{:best, best}] -> heat_loss > best
              end

            if not worse_than_best do
              if straight_count < max_straight do
                simulate_movement(
                  map,
                  goal,
                  min_straight,
                  max_straight,
                  apply_direction(position, direction),
                  direction,
                  straight_count + 1,
                  heat_loss,
                  Map.put(visited, position, true)
                )
              end

              if straight_count >= min_straight do
                new_direction = turn(direction, :left)

                simulate_movement(
                  map,
                  goal,
                  min_straight,
                  max_straight,
                  apply_direction(position, new_direction),
                  new_direction,
                  1,
                  heat_loss,
                  Map.put(visited, position, true)
                )

                new_direction = turn(direction, :right)

                simulate_movement(
                  map,
                  goal,
                  min_straight,
                  max_straight,
                  apply_direction(position, new_direction),
                  new_direction,
                  1,
                  heat_loss,
                  Map.put(visited, position, true)
                )
              end
            end
          end
        end

      _ ->
        :ok
    end
  end

  def part1(map) do
    :ets.new(:movement_memo, [:named_table])
    :ets.new(:best_memo, [:named_table])

    goal = Enum.map(map, &elem(&1, 0)) |> Enum.max()
    simulate_movement(map, goal, 1, 3)

    best =
      :ets.tab2list(:movement_memo)
      |> Enum.filter(fn {{position, _, _}, _} -> position == goal end)
      |> Enum.map(fn {_, heat_loss} -> heat_loss end)
      |> Enum.min()

    :ets.delete(:movement_memo)
    :ets.delete(:best_memo)

    best
  end

  def part2(map) do
    :ets.new(:movement_memo, [:named_table])
    :ets.new(:best_memo, [:named_table])

    goal = Enum.map(map, &elem(&1, 0)) |> Enum.max()
    simulate_movement(map, goal, 4, 10)

    best =
      :ets.tab2list(:movement_memo)
      |> Enum.filter(fn {{position, _, _}, _} -> position == goal end)
      |> Enum.map(fn {_, heat_loss} -> heat_loss end)
      |> Enum.min()

    :ets.delete(:movement_memo)
    :ets.delete(:best_memo)

    best
  end
end