import java.lang.Long.numberOfTrailingZeros
import java.util.*
import kotlin.math.abs
import kotlin.math.min
import kotlin.math.sqrt

fun calcCanonicalPrimesSieveOfEratosthenes(n: Int): List<Int> {
    return primeSequence(n).takeWhile { it < n }.toList()
}

fun sieveIsPrimeFunction(n: Int): (Int) -> Boolean = sieveIsPrimeFunction(sieveOfErastosthenes(n))

fun sieveIsPrimeFunction(bits: BitSet): (Int) -> Boolean = { i -> (i == 2) || (((i and 1) == 1) && !bits[i ushr 1]) }

fun primeSequence(n: Int) = primeSequence(sieveOfErastosthenes(n))

fun primeSequence(bits: BitSet) = generateSequence(2) { (bits.nextClearBit((it + 1) ushr 1) shl 1) + 1 }

fun sieveOfErastosthenes(n: Int): BitSet {
    val nHalf = (n ushr 1) + 1
    val bits = BitSet(nHalf)
    bits.set(0)
    val upperLimit = sqrt(n.toDouble()).toInt()
    for (i in 3..upperLimit step 2) {
        if (!bits[i ushr 1]) {
            var v = (i * i) ushr 1
            while (v < nHalf) {
                bits.set(v)
                v += i
            }
        }
    }
    return bits
}

fun primeFactors(n: Long, bits: BitSet): List<Long> {
    if (n < bits.length()) {
        if (sieveIsPrimeFunction(bits).invoke(n.toInt())) {
            return listOf(n)
        }
        return primeFactorsTrialDivision(n.toInt(), bits).map(Int::toLong)
    }
    if (n.toBigInteger().isProbablePrime(20)) {
        return listOf(n)
    }
    val (trivialFactors, remainder) = partialPrimeFactorsTrialDivision(n, 5000, bits)
    if (remainder == 1L) {
        return trivialFactors
    }
    val factors = ArrayList(trivialFactors)
    if (!remainder.toBigInteger().isProbablePrime(20)) {
        val divisor = rhoBrent(remainder)
        factors.addAll(primeFactors(divisor, bits))
        val newRemainder = remainder / divisor
        if (newRemainder > 1) {
            factors.addAll(primeFactors(newRemainder, bits))
        }
    } else {
        factors.add(remainder)
    }
    return factors.sorted()
}

private fun primeFactorsTrialDivision(n: Int, bits: BitSet): List<Int> {
    val factors = ArrayList<Int>()
    var sq = Math.sqrt(n.toDouble()).toInt()
    var r = n
    for (prime in primeSequence(bits)) {
        if (prime > sq) {
            if (r > 1) {
                factors.add(r)
            }
            break
        }
        if ((r % prime) == 0) {
            r /= prime
            factors.add(prime)
            while ((r % prime) == 0) {
                r /= prime
                factors.add(prime)
            }
            sq = sqrt(n.toDouble()).toInt()
            if (r == 1) {
                break
            }
        }
    }
    return factors
}

private fun partialPrimeFactorsTrialDivision(n: Long, maxPrime: Int, bits: BitSet): Pair<List<Long>, Long> {
    val factors = ArrayList<Long>()
    var sq = sqrt(n.toDouble()).toLong()
    var r = n
    for (prime in primeSequence(bits)) {
        if (prime > sq) {
            if (r > 1) {
                factors.add(r)
                r = 1L
            }
            break
        }
        if (prime > maxPrime) {
            break
        }
        if ((r % prime) == 0L) {
            r /= prime
            factors.add(prime.toLong())
            while ((r % prime) == 0L) {
                r /= prime
                factors.add(prime.toLong())
            }
            sq = sqrt(n.toDouble()).toLong()
            if (r == 1L) {
                break
            }
        }
    }
    return factors to r
}

private fun rhoBrent(n: Long): Long {
    val x0 = 2L
    val m = 25
    var cst = 1_000_000_007L
    var y = x0
    var r = 1
    do {
        val x = y
        for (i in 0 until r) {
            val y2 = y * y
            y = ((y2 + cst) % n)
        }
        var k = 0
        do {
            val bound = min(m, r - k)
            var q = 1L
            for (i in -3 until bound) { //start at -3 to ensure we enter this loop at least 3 times
                val y2 = y * y
                y = ((y2 + cst) % n)
                val divisor = abs(x - y)
                if (0L == divisor) {
                    cst += 1_000_000_007L
                    k = -m
                    y = x0
                    r = 1
                    break
                }
                val prod = divisor * q
                q = prod % n
                if (q == 0L) {
                    return gcdPositive(abs(divisor), n)
                }
            }
            val out = gcdPositive(abs(q), n)
            if (out != 1L) {
                return out
            }
            k += m
        } while (k < r)
        r += r
    } while (true)
}

fun gcdPositive(aIn: Long, bIn: Long): Long {
    if (aIn == 0L) {
        return bIn
    } else if (bIn == 0L) {
        return aIn
    }

    var a = aIn
    var b = bIn
    val aTwos = numberOfTrailingZeros(a)
    a = a ushr aTwos
    val bTwos = numberOfTrailingZeros(b)
    b = b ushr bTwos
    val shift = min(aTwos, bTwos)

    while (a != b) {
        val delta = a - b
        b = min(a, b)
        a = abs(delta)
        a = a ushr numberOfTrailingZeros(a)
    }

    return a shl shift
}

// ax + by = gcdExtendedPositive(a, b)
fun extendedGcd(a: Long, b: Long): Pair<Long, Pair<Long, Long>> {
    var old_r = a
    var r = b
    var old_s = 1L
    var s = 0L
    var old_t = 0L
    var t = 1L
    while (r != 0L) {
        val q = old_r / r
        val rtmp = old_r
        old_r = r
        r = rtmp - q * r
        val stmp = old_s
        old_s = s
        s = stmp - q * s
        val ttmp = old_t
        old_t = t
        t = ttmp - q * t
    }

    return old_r to (old_s to old_t)
}

fun extendedGcd(v: List<Long>): Pair<Long, List<Long>> {
    if (v.size < 2) throw IllegalArgumentException("Expected at least 2 elements")

    val gcds = ArrayList<Long>(v.size)
    val coeffs = ArrayList<Long>(v.size)
    var (gcd, p1) = extendedGcd(v[0], v[1])
    coeffs.add(p1.first)
    coeffs.add(p1.second)
    gcds.add(gcd)
    gcds.add(gcd)
    for (i in 2 until v.size) {
        val (gcdnew, pi) = extendedGcd(gcd, v[i])
        gcd = gcdnew
        coeffs.add(pi.second)
        gcds.add(gcd)
    }
    for (i in gcds.indices) {
        if (gcds[i] != gcd) {
            coeffs[i] *= gcds[i] / gcd
        }
    }
    return gcd to coeffs
}

fun calcPrimeFactorsAndPhi(
    n: Long,
    primes: MutableList<Long>,
    allPrimes: MutableSet<Long>
): Pair<List<Pair<Long, Int>>, Long> {
    val factors = ArrayList<Pair<Long, Int>>()
    var phi = 1L
    var rem = n
    if (rem and 1 == 0L) {
        var times = 1
        rem = rem ushr 1
        while (rem and 1 == 0L) {
            times++
            rem = rem ushr 1
            phi *= 2
        }
        factors.add(2L to times)
    }
    if (rem > 1) {
        if (!allPrimes.contains(rem)) {
            var primePos = 1
            var prime = 3L
            var sq = Math.sqrt(rem.toDouble()).toInt()
            while (prime <= sq) {
                if (rem % prime == 0L) {
                    var times = 1
                    rem /= prime
                    phi *= prime - 1
                    while (rem % prime == 0L) {
                        times++
                        rem /= prime
                        phi *= prime
                    }
                    factors.add(prime to times)
                    if (rem == 1L) {
                        break
                    }
                    sq = Math.sqrt(rem.toDouble()).toInt()
                    if (primePos > primes.lastIndex) {
                        primes.add(prime)
                        allPrimes.add(prime)
                    }
                } else if ((primePos > primes.lastIndex) && isPrime(prime, sq, primes)) {
                    primes.add(prime)
                    allPrimes.add(prime)
                    primePos++
                }
                if (primePos > primes.lastIndex) {
                    prime += 2
                } else {
                    prime = primes[primePos++]
                }
            }
        }
        if (rem > 1) {
            factors.add(rem to 1)
            allPrimes.add(rem)
            phi *= (rem - 1)
        }
    }
    return factors to phi
}

private fun isPrime(i: Long, sq: Int, primes: Collection<Long>): Boolean {
    for (p in primes) {
        if (p > sq) {
            break
        }
        if (i % p == 0L) {
            return false
        }
    }
    return true
}