onLights =
  File.stream!("inputs/day18.txt")
  |> Enum.with_index()
  |> Enum.flat_map(fn {line, y} ->
    String.trim(line)
    |> String.graphemes()
    |> Enum.with_index()
    |> Enum.filter(fn {on, _x} -> on == "#" end)
    |> Enum.map(fn {_, x} -> {x, y} end)
  end)
  |> MapSet.new()

lightStream = fn initial, stuckLights ->
  Stream.iterate(initial, fn onLights ->
    newOnLights =
      stuckLights ++
        Enum.flat_map(0..99, fn y ->
          Enum.filter(0..99, fn x ->
            currentlyOn = MapSet.member?(onLights, {x, y})

            neighborCount =
              Enum.filter(
                [
                  MapSet.member?(onLights, {x - 1, y - 1}),
                  MapSet.member?(onLights, {x, y - 1}),
                  MapSet.member?(onLights, {x + 1, y - 1}),
                  MapSet.member?(onLights, {x - 1, y}),
                  MapSet.member?(onLights, {x + 1, y}),
                  MapSet.member?(onLights, {x - 1, y + 1}),
                  MapSet.member?(onLights, {x, y + 1}),
                  MapSet.member?(onLights, {x + 1, y + 1})
                ],
                & &1
              )
              |> Enum.count()

            neighborCount == 3 || (currentlyOn && neighborCount == 2)
          end)
          |> Enum.map(fn x -> {x, y} end)
        end)

    MapSet.new(newOnLights)
  end)
end

[after100] = Stream.drop(lightStream.(onLights, []), 100) |> Enum.take(1)
IO.puts("Lights on: #{MapSet.size(after100)}")

[after100Stuck] =
  Stream.drop(lightStream.(onLights, [{0, 0}, {99, 0}, {0, 99}, {99, 99}]), 100) |> Enum.take(1)

IO.puts("Lights on with the corners stuck: #{MapSet.size(after100Stuck)}")