import numpy as np

# First 50 non-trivial zeta zeros (positive imaginary parts)
# Zeta zeros come in conjugate pairs rho = 1/2 ± i*gamma
ZETA_GAMMAS = [
    14.1347251417, 21.0220396388, 25.0108575801, 30.4248761259,
    32.9350615877, 37.5861781588, 40.9187190121, 43.3270732809,
    48.0051508812, 49.7738324777, 52.9703214777, 56.4462476971,
    59.3470440026, 60.8317785246, 65.1125440481, 67.0798105295,
    69.5464017112, 72.0671576745, 75.7046906991, 77.1448400689,
    79.3373750205, 82.9103808541, 84.7354929805, 87.4252746131,
    88.8091112076, 92.4918992706, 94.6513440405, 95.8706342582,
    98.8311942182, 101.3178510057, 103.7255380405, 105.4466230523,
    107.1686111842, 111.0295355432, 111.8746591769, 114.3202209154,
    116.2266803209, 117.7896968966, 121.3708050022, 122.9468292935,
    124.2568185544, 127.5166838796, 129.5787041981, 131.0876881309,
    133.4977372027, 134.7565457538, 138.1160420545, 139.7362089521,
    141.1237074040, 143.1118458076
]

def mobius(n):
    """Compute the Möbius function."""
    if n == 1:
        return 1
    count = 0
    p = 2
    while p * p <= n:
        if n % p == 0:
            if n % (p * p) == 0:
                return 0
            n //= p
            count += 1
        p += 1 if p == 2 else 2
    if n > 1:
        count += 1
    return -1 if count % 2 else 1

def psi_explicit(x, K=50):
    """
    Chebyshev psi(x) via the explicit formula.
    
    psi(x) = x - sum over zeta zeros of x^rho/rho - log(2*pi) - 0.5*log(1 - x^(-2))
    """
    if x <= 1:
        return 0.0
    
    result = x - np.log(2 * np.pi)
    
    # Sum over conjugate pairs rho = 1/2 ± i*gamma
    # Each pair contributes: 2 * Re[x^rho / rho]
    for gamma in ZETA_GAMMAS[:K]:
        rho = 0.5 + 1j * gamma
        term = x**rho / rho
        result -= 2 * np.real(term)
    
    # Trivial zeros contribution
    if x > 1:
        result -= 0.5 * np.log(1 - x**(-2))
    
    return result

def theta_from_psi(x, K=50):
    """
    First Chebyshev theta(x) via Möbius inversion:
    
    psi(x) = sum_{k=1}^inf theta(x^{1/k})
    => theta(x) = sum_{k=1}^inf mu(k) * psi(x^{1/k})
    """
    if x <= 1:
        return 0.0
    
    total = 0.0
    k_max = int(np.floor(np.log2(x))) + 1
    
    for k in range(1, k_max + 1):
        mu = mobius(k)
        if mu == 0:
            continue
        total += mu * psi_explicit(x**(1.0 / k), K)
    
    return total

def pi_from_theta(x, K=50):
    """
    Count primes up to x by detecting jumps in theta(x).
    theta jumps by log(p) at each prime p.
    """
    if x < 2:
        return 0
    
    n_max = int(np.floor(x))
    count = 0
    
    for n in range(2, n_max + 1):
        theta_n = theta_from_psi(float(n), K)
        theta_prev = theta_from_psi(float(n - 1), K)
        jump = theta_n - theta_prev
        
        # If the jump is close to log(n), n is prime
        if jump / np.log(n) > 0.5:
            count += 1
    
    return count

def nth_prime_explicit(n, K=50):
    """
    Find the nth prime by scanning upward and counting primes.
    """
    #if n < 1:
    #    raise ValueError("n must be a positive integer")
    
    count = 0
    x = 1.0
    
    while count < n:
        x += 1.0
        theta_n = theta_from_psi(x, K)
        theta_prev = theta_from_psi(x - 1, K)
        jump = theta_n - theta_prev
        
        if jump / np.log(x) > 0.5:
            count += 1
    
    return int(x)

if __name__ == "__main__":
    known_primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29,
                    31, 37, 41, 43, 47, 53, 59, 61, 67, 71]
    
    print("n\tapprox p(n)\texact p(n)\terror")
    for n in range(1, 21):
        approx = nth_prime_explicit(n, K=50)
        exact = known_primes[n - 1]
        print(f"{n}\t{approx}\t\t{exact}\t\t{approx - exact}")