require "./euler"

module Euler
  class Prime
    include Iterator(NumType | BigInt)

    def initialize()
      @sieve_size = 16
      @sieved_up_to = 0

      @primes = Array(Bool).new(@sieve_size, true)
      @primes[0] = false
      @primes[1] = false

      @index = 0
      iterative_eratosthenes_sieve
    end

    private def check_and_increase_sieve
      if @index > @primes.size - 1
        @sieve_size *= 2
        until @primes.size == @sieve_size
          @primes << true
        end
      end

      iterative_eratosthenes_sieve
    end

    private def iterative_eratosthenes_sieve
      2.step(to: Math.sqrt(@sieve_size)) do |i|
        if @primes[i]
          (i * i >= @sieved_up_to ? i * i : @sieved_up_to + i - (@sieved_up_to % i)).step(to: @sieve_size-1, by: i) do |j|
            @primes[j] = false
          end
        end
      end
      @sieved_up_to = @sieve_size-1
    end

    def next
      @index += 1
      check_and_increase_sieve

      until @primes[@index]
        @index += 1
        check_and_increase_sieve
      end
      @index
    end

    def self.eratosthenes_sieve(n)
      sieve = Array.new(n + 1, true)
      sieve[0] = false
      sieve[1] = false

      2.step(to: Math.sqrt(n)) do |i|
        if sieve[i]
          (i * i).step(to: n, by: i) do |j|
            sieve[j] = false
          end
        end
      end

      result = sieve.map_with_index { |b, i| b ? i : nil }.compact
    end

    macro trial_division_method(given_type)
      def self.trial_division(n : {{given_type}})
        factors = [] of {{given_type}}
        check = ->(p: {{given_type}}) {
          q, r = n.divmod(p)
          while r.zero?
            factors << p
            n = q
            q, r = n.divmod(p)
          end
        }

        check.call(2)
        check.call(3)
        p = 5
        while p * p <= n
          check.call(p)
          p += 2
          check.call(p)
          p += 4
        end
        factors << n if n > 1
        factors
      end
    end

    trial_division_method(NumType)
    trial_division_method(BigInt)

    def self.prime_factorization(n : NumType | BigInt)
      result = {} of (NumType | BigInt) => Int32
      factors = self.trial_division(n)

      factors.each do |f|
        result[f] = 0
        num = n
        while num % f == 0
          result[f] += 1
          num /= f
        end
      end

      result
    end
  end
end