from lib_prime import primes
from math import log, ceil


def c(n_base: int) -> int:
    """Compute the number of prime pairs (p, q) where
    p**q * q**p <= n_base**n_base and p < q"""
    digits = ceil(log(n_base, 10)) + 1
    n = n_base**n_base
    n_log = log(n)
    ps = primes(2 * 10**digits)
    p, q = ps[0], ps[-1]
    assert (p**q * q**p) > n, "largest prime too low"
    j_max = len(ps) - 1

    t = 0
    for i in range(len(ps)):
        p, q = ps[i], ps[i + 1]

        # Crux:
        #             (p**q * q**p) > n
        # <=>       ln(p**q * q**p) > ln(n)
        # <=>   ln(p**q) + ln(q**p) > ln(n)
        # <=> q * ln(p) + p * ln(q) > ln(n)
        # Much lower numbers and therefore faster to compute.
        if q * log(p) + p * log(q) > n_log:
            break

        # Binary search to find the largest j for which p**q * q**p <= n:
        j_lo, j_hi = i + 1, j_max
        while (j_hi - j_lo) > 1:
            j_mid = (j_hi + j_lo) // 2
            p, q = ps[i], ps[j_mid]
            if q * log(p) + p * log(q) > n_log:
                j_hi = j_mid
            else:
                j_lo = j_mid
        j_max = j_lo + 1

        # All primes lower than ps[j_lo] have to be counted.
        t += j_lo - i

        # p_ref, q_ref = ps[i], ps[j_lo]
        # print(p_ref, q_ref, "<=")

        # Turns out the binary search is not necessary with the optimized
        # version of the comparison and the following works just as fine.
        # for j in range(j_max, i, -1):
        #     p, q = ps[i], ps[j]
        #     if q * log(p) + p * log(q) > n_log:
        #         j_max = j
        #     else:
        #         assert p == p_ref
        #         assert q == q_ref
        #         t += (j - i)
        #         break
    return t


def euler_800():
    assert c(800) == 10790
    return c(800800)


if __name__ == "__main__":
    solution = euler_800()
    print(f"e800.py: {solution}")
    assert(solution == 1412403576)