defmodule Prime do
  @doc """
  Generates the nth prime.
  """
  @spec nth(non_neg_integer) :: non_neg_integer
  def nth(count) when count >= 1 do
    0
    |> Stream.iterate(&(&1 + 1))
    |> Stream.filter(&prime?/1)
    |> Stream.take(count + 1)
    |> Enum.to_list()
    |> List.last()
  end

  @spec prime?(non_neg_integer) :: boolean
  defp prime?(number) when number in 1..3, do: true
  defp prime?(number) when rem(number, 2) == 0, do: false
  defp prime?(number) when number > 3, do: check_prime(number)

  defp check_prime(number, prime \\ 3)
  defp check_prime(number, number), do: true

  defp check_prime(number, prime) do
    if prime >= trunc(number ** 0.5) + 1,
      do: true,
      else: rem(number, prime) != 0 && check_prime(number, prime + 2)
  end

  # defp check_prime(number, prime), do: rem(number, prime) != 0 && check_prime(number, prime + 2)
end