defmodule Day6 do
  def run do
    map =
      File.read!("data.txt")
      |> String.split("\n")
      |> Enum.drop(-1)
      |> Enum.map(&String.graphemes(&1))

    height = length(map)
    width = length(hd(map))

    # find initial location
    idx = Enum.concat(map) |> Enum.find_index(&(&1 != "#" and &1 != "."))
    location = {div(idx, height), rem(idx, width)}

    initialState = %{
      map: map,
      visited: MapSet.new(),
      location: location,
      bearing: 0,
      height: height,
      width: width
    }

    visited = walk(initialState, false)
    IO.puts("Part 1: #{MapSet.size(visited)}")

    # for each visited location, try putting an obstacle there
    num_positions =
      Enum.reduce(visited, 0, fn {row, col}, acc ->
        # generate new map with obstacle in the path
        new_row = List.replace_at(Enum.at(map, row), col, "#")
        new_map = List.replace_at(map, row, new_row)
        acc + walk(%{initialState | map: new_map}, true)
      end)

    IO.puts("Part 2: #{num_positions}")
  end

  defp walk(state, detect_loops?) do
    # if we're detecting loops, and this location & bearing is already in our map, get out
    if detect_loops? and MapSet.member?(state.visited, {state.location, state.bearing}) do
      1
    else
      # add current location to set
      visited =
        if detect_loops? do
          # if we're detecting loops, also record the bearing
          MapSet.put(state.visited, {state.location, state.bearing})
        else
          MapSet.put(state.visited, state.location)
        end

      {row, col} = state.location

      # grab next location
      {next_row, next_col} =
        case state.bearing do
          0 -> {row - 1, col}
          1 -> {row, col + 1}
          2 -> {row + 1, col}
          3 -> {row, col - 1}
        end

      in_bounds =
        next_row >= 0 and next_row < state.height and
          next_col >= 0 and next_col < state.width

      if !in_bounds do
        # if not in bounds, end recursion
        if detect_loops? do
          # if we're detecting loops, we didn't find one. return 0
          0
        else
          # otherwise return the list of visited locations
          visited
        end
      else
        cell = Enum.at(state.map, next_row) |> Enum.at(next_col)

        new_state =
          if cell == "#" do
            # keep location the same, turn 90° right
            new_bearing = Integer.mod(state.bearing + 1, 4)
            %{state | visited: visited, bearing: new_bearing}
          else
            %{state | visited: visited, location: {next_row, next_col}}
          end

        walk(new_state, detect_loops?)
      end
    end
  end
end

Day6.run()