defmodule Aoc2023.Day16 do
  use Aoc2023

  defp emissions(char, dir) do
    case {char, dir} do
      {?., dir} -> [dir]
      {?-, {0, dir}} -> [{0, dir}]
      {?-, {_, 0}} -> [{0, -1}, {0, 1}]
      {?|, {dir, 0}} -> [{dir, 0}]
      {?|, {0, _}} -> [{-1, 0}, {1, 0}]
      {?/, {dir, 0}} -> [{0, -dir}]
      {?/, {0, dir}} -> [{-dir, 0}]
      {?\\, {dir, 0}} -> [{0, dir}]
      {?\\, {0, dir}} -> [{dir, 0}]
    end
  end

  def parse(input) do
    input
    |> lines
    |> Enum.map(&String.to_charlist/1)
    |> Enum.with_index()
    |> Enum.flat_map(fn {row, y} ->
      row |> Enum.with_index() |> Enum.map(fn {c, x} -> {{y, x}, c} end)
    end)
    |> Enum.into(%{})
  end

  defp run(input, heads) do
    loop = fn recur, energized, heads ->
      energized_ = MapSet.union(energized, heads |> MapSet.new())

      heads_ =
        heads
        |> Enum.flat_map(fn {coord, dir} ->
          char = Map.get(input, coord)

          if char,
            do:
              emissions(char, dir)
              |> Enum.map(&{add_vec(coord, &1), &1})
              |> Enum.reject(&MapSet.member?(energized_, &1))
              |> Enum.filter(&Map.has_key?(input, elem(&1, 0))),
            else: []
        end)

      if length(heads_) > 0, do: recur.(recur, energized_, heads_), else: energized_
    end

    loop.(loop, MapSet.new(), heads) |> Enum.map(pipe(elem(0))) |> MapSet.new() |> MapSet.size()
  end

  def part1(input) do
    run(input, [{{0, 0}, {0, 1}}])
  end

  def part2(input) do
    {my, mx} = input |> Enum.map(pipe(elem(0))) |> Enum.max()

    [
      0..my |> Enum.map(&{{&1, 0}, {0, 1}}),
      0..my |> Enum.map(&{{&1, mx}, {0, -1}}),
      0..mx |> Enum.map(&{{0, &1}, {1, 0}}),
      0..mx |> Enum.map(&{{my, &1}, {-1, 0}})
    ] |> Enum.concat |> Enum.map(&run(input, [&1])) |> Enum.max
  end
end