import numpy as np
from scipy.special import expi
from scipy.optimize import bisect

# First 20 non-trivial zeta zeros (imaginary parts)
# For better accuracy, you can extend this list (e.g., from OEIS A002410)
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
]

def li(x):
    """
    Logarithmic integral li(x) = ∫₀ˣ dt / ln t.
    For real x > 1, li(x) = Ei(ln x).  For complex, we use expi.
    For x <= 1, we return a real value via analytic continuation,
    but in practice we only call with x > 1 or complex x with magnitude > 1.
    """
    if isinstance(x, complex) or x > 1:
        # expi handles complex arguments
        return np.real(expi(np.log(x + 0j)))  # cast to complex to be safe
    else:
        # For 0 < x <= 1, li(x) is defined as the Cauchy principal value,
        # but we rarely need it; return 0 as a fallback.
        return 0.0

def pi_explicit(x, K=10):
    """
    Approximate π(x) using the Riemann explicit formula with K non-trivial zeros.
    """
    if x <= 1:
        return 0.0

    result = li(x) - np.log(2)   # -log 2 term

    # Trivial zeros contribution: 0.5 * log(1 - x^{-2})
    if x > 1:
        result += float(np.real(0.5 * np.log(1 - x**(-2))))

    # Non-trivial zeros: sum over rho = 1/2 + i*gamma
    for gamma in ZETA_GAMMAS[:K]:
        rho = 0.5 + 1j * gamma
        result -= np.real(li(x ** rho))

    # The integral term ∫_x^∞ dt/(t(t²-1)log t) is small for large x;
    # we omit it for simplicity, but for very accurate results it could be added.
    return result

def nth_prime_explicit(n, K=10, tol=1e-3, max_iter=100):
    """
    Compute the nth prime p(n) by numerically solving π(p(n)) = n.
    Uses the explicit formula for π(x) with K zeta zeros.
    """
    if n < 1:
        raise ValueError("n must be a positive integer")
    if n == 1:
        return 2
    if n == 2:
        return 3

    # Initial bounds: using the asymptotic p(n) ~ n log n, but we need a bracket.
    # We use the known bounds: n (log n + log log n - 1) < p(n) < n (log n + log log n)
    # for n >= 6. For smaller n we handle separately.
    def lower_bound(n):
        if n < 6:
            return max(2, n)  # just a safe fallback
        logn = np.log(n)
        loglogn = np.log(logn)
        return n * (logn + loglogn - 1)

    def upper_bound(n):
        if n < 6:
            return max(10, n * 2)
        logn = np.log(n)
        loglogn = np.log(logn)
        return n * (logn + loglogn)

    a = max(2.0, lower_bound(n))
    b = upper_bound(n) + 10.0   # add some margin

    # Ensure π(a) < n and π(b) > n
    while pi_explicit(a, K) >= n:
        a /= 1.1
    while pi_explicit(b, K) < n:
        b *= 1.1

    # Use bisection (or scipy.optimize.bisect)
    def f(x):
        return pi_explicit(x, K) - n

    try:
        root = bisect(f, a, b, xtol=tol)
        return root
    except ValueError:
        # Fallback: use a grid search if bisection fails
        x = np.linspace(a, b, 1000)
        pi_vals = [pi_explicit(xi, K) for xi in x]
        idx = np.argmin(np.abs(np.array(pi_vals) - n))
        return x[idx]

# Example usage
if __name__ == "__main__":
    # Test with known primes
    known_primes = {1:2, 2:3, 3:5, 4:7, 5:11, 6:13, 7:17, 8:19, 9:23, 10:29,
                    11:31, 12:37, 13:41, 14:43, 15:47, 16:53, 17:59, 18:61, 19:67, 20:71}
    print("n\tapprox p(n)\texact p(n)\terror")
    for n in range(1, 21):
        approx = nth_prime_explicit(n, K=10)
        exact = known_primes[n]
        print(f"{n}\t{approx:.3f}\t\t{exact}\t\t{approx-exact:.3f}")