/*
Copyright © 2025 filifa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package sieve

// NOTE: these formulas come from https://web.archive.org/web/20151224013638/http://www.maa.org/sites/default/files/Walter_D22068._Stangl.pdf

func updatePowersOfTwoCoprime(sieve []uint, n uint) {
	for q := uint(8); 2*q < n; q *= 2 {
		// q(2^n) = 2^(n-3) for n >= 3
		sieve[2*q] = 2 * sieve[q]
	}
}

func updatePowersOfOddPrimesCoprime(sieve []uint, p uint, n uint) {
	for q := p; p*q < n; q *= p {
		// q(p^n) = (p^n - p^(n-1)) / 2
		sieve[p*q] = (p*q - q) / 2
	}
}

func updatePowersOfTwo(sieve []uint, n uint) {
	k := 1
	for q := uint(1); 2*q < n; q *= 2 {
		if k%2 == 0 {
			// s(2^n) = (2^(n-1) + 4) / 3 for even n
			sieve[2*q] = (q + 4) / 3
		} else {
			// s(2^n) = (2^(n-1) + 5) / 3 for odd n
			sieve[2*q] = (q + 5) / 3
		}

		k += 1
	}
}

func updatePowersOfOddPrimes(sieve []uint, p uint, n uint) {
	k := 2
	for q := p; p*q < n; q *= p {
		if p == q {
			// s(p^2) = (p^2 - p + 2) / 2
			sieve[p*q] = (p*p - p + 2) / 2
		} else if k%2 == 0 {
			// s(p^n) = (p^(n+1) + p + 2) / (2*(p+1)) for even n
			sieve[p*q] = (p*p*q + p + 2) / (2 * (p + 1))
		} else {
			// s(p^n) = (p^(n+1) + 2*p + 1) / (2*(p+1)) for odd n
			sieve[p*q] = (p*p*q + 2*p + 1) / (2 * (p + 1))
		}

		k += 1
	}
}

/*
QuadraticResidues computes the number of quadratic residues modulo k for k=1 to n.

see https://oeis.org/A046073
*/
func QuadraticResidues(n uint, coprime bool, buflen uint) chan uint {
	sieve := make([]uint, n)
	for i := uint(0); i < n; i++ {
		sieve[i] = 1
	}

	ch := make(chan uint, buflen)
	go func() {
		defer close(ch)
		if coprime {
			sieveQRCoprime(sieve, n, ch)
		} else {
			sieveQR(sieve, n, ch)
		}
	}()

	return ch
}

func sieveQRCoprime(sieve []uint, n uint, ch chan uint) {
	for i := uint(0); i < n; i++ {
		if i == 0 || i == 1 || i == 4 || i == 6 || i == 8 || i == 12 || i == 24 || sieve[i] != 1 {
			ch <- sieve[i]
			continue
		}

		if i == 2 {
			updatePowersOfTwoCoprime(sieve, n)
		} else {
			sieve[i] = (i - 1) / 2
			updatePowersOfOddPrimesCoprime(sieve, i, n)
		}

		updateMultiples(sieve, i, n, false)
		ch <- sieve[i]
	}
}

func sieveQR(sieve []uint, n uint, ch chan uint) {
	for i := uint(0); i < n; i++ {
		if i == 0 || i == 1 || sieve[i] != 1 {
			ch <- sieve[i]
			continue
		}

		if i == 2 {
			updatePowersOfTwo(sieve, n)
		} else {
			sieve[i] = (i + 1) / 2
			updatePowersOfOddPrimes(sieve, i, n)
		}

		updateMultiples(sieve, i, n, false)
		ch <- sieve[i]
	}
}