// Package erastothenes implements a sieve of Erastothenes with various optimizations.
package erastothenes

import (
	"primes"
	"sync"
)

// Unoptimized is the basic method without any optimization or thinking.
func Unoptimized(n int) []bool {
	sieve := make([]bool, n+1)
	for i := 2; i <= n; i++ {
		for j := 2 * i; j <= n; j += i {
			sieve[j] = true
		}
	}
	return sieve
}

// SkipEven skips even numbers.
func SkipEven(n int) []bool {
	sieve := make([]bool, n+1)
	for i := 2; i <= n; i += 2 {
		for j := 2 * i; j <= n; j += i {
			sieve[j] = true
		}
		if i == 2 {
			i--
		}
	}
	return sieve
}

// StartPower starts with j = i*i.
func StartPower(n int) []bool {
	sieve := make([]bool, n+1)
	for i := 2; i <= n; i++ {
		for j := i * i; j <= n; j += i {
			sieve[j] = true
		}
	}
	return sieve
}

// OnlyTilRoot stops outer loop at i*i<n.
func OnlyTilRoot(n int) []bool {
	sieve := make([]bool, n+1)
	for i := 2; i*i <= n; i++ {
		for j := 2 * i; j <= n; j += i {
			sieve[j] = true
		}
	}
	return sieve
}

// Optimized has all the optimizations above.
func Optimized(n int) []bool {
	sieve := make([]bool, n+1)
	for i := 2; i*i <= n; i += 2 {
		for j := i * i; j <= n; j += i {
			sieve[j] = true
		}
		if i == 2 {
			i--
		}
	}
	return sieve
}

// WaitGroup based parallel sieve.
func WaitGroup(limit int) []bool {
	sieve := make([]bool, limit+1)
	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		// handle even numbers
		for i := 2 * 2; i <= limit; i += 2 {
			sieve[i] = true
		}
		wg.Done()
	}()
	for p := 3; p*p <= limit; p += 2 {
		wg.Add(1)
		go func(p int) {
			for i := p * p; i <= limit; i += p {
				sieve[i] = true
			}
			wg.Done()
		}(p)
	}
	wg.Wait()
	return sieve
}

func Nth(sieve string, limit int) int {
	if limit < 1 {
		return 0
	}
	var s []bool
	switch sieve {
	case "unoptimized":
		s = Unoptimized(primes.MAX)
	case "skipeven":
		s = SkipEven(primes.MAX)
	case "power":
		s = StartPower(primes.MAX)
	case "root":
		s = OnlyTilRoot(primes.MAX)
	case "opt":
		s = Optimized(primes.MAX)
	case "waitgroup":
		s = WaitGroup(primes.MAX)
	}

	i := 2
	for count := 0; count < limit; i++ {
		if !s[i] {
			count++
		}
	}
	return i - 1
}