; counting primes

(define (isqrt n)
  (if (not (and (positive? n) (integer? n)))
      (error 'isqrt "must be positive integer")
      (let loop ((x n))
        (let ((y (quotient (+ x (quotient n x)) 2)))
          (if (< y x) (loop y) x)))))

(define (primes n)
  (let* ((max-index (quotient (- n 3) 2))
         (v (make-vector (+ 1 max-index) #t)))
    (let loop ((i 0) (ps '(2)))
      (let ((p (+ i i 3)) (startj (+ (* 2 i i) (* 6 i) 3)))
        (cond ((>= (* p p) n)
               (let loop ((j i) (ps ps))
                  (cond ((> j max-index) (reverse ps))
                        ((vector-ref v j)
                          (loop (+ j 1) (cons (+ j j 3) ps)))
                        (else (loop (+ j 1) ps)))))
              ((vector-ref v i)
                (let loop ((j startj))
                  (if (<= j max-index)
                      (begin (vector-set! v j #f)
                             (loop (+ j p)))))
                      (loop (+ 1 i) (cons p ps)))
              (else (loop (+ 1 i) ps)))))))

(define (make-pi-vector)
  (with-output-to-file "pi-vector.ss" (lambda ()
    (display "#(") (newline) (display 0) (newline)
    (let* ((ps (cdr (primes (+ (isqrt #e1e12) 1))))
           (qs (map (lambda (p) (modulo (* -1/2 (+ #e1e8 1 p)) p)) ps))
           (z (length (primes #e1e8))))
      (do ((t #e1e8 (+ t #e1e5 #e1e5))
           (qs qs (map (lambda (p q) (modulo (- q #e1e5) p)) ps qs)))
          ((> t #e1e12) (display ")") (newline))
        (when (zero? (modulo t #e1e8)) (display z) (newline))
        (let ((bv (make-vector #e1e5 #t)))
          (do ((qs qs (cdr qs)) (ps ps (cdr ps))) ((null? qs))
            (do ((j (car qs) (+ j (car ps)))) ((<= #e1e5 j))
              (vector-set! bv j #f)))
          (do ((j 0 (+ j 1))) ((= j #e1e5))
            (if (vector-ref bv j) (set! z (+ z 1))))))))))

(define (prime-pi n)
  (if (< #e1e12 n) (error 'prime-pi "out of range")
    (if (< n #e1e8) (length (primes n))
      (let* ((n (if (even? n) (- n 1) n))
             (b (quotient n #e1e8)) (l (* b #e1e8))
             (ps (cdr (primes (+ (isqrt n) 1))))
             (qs (map (lambda (p) (modulo (* -1/2 (+ l 1 p)) p)) ps))
             (z (vector-ref pi-vector b)) (pi #f))
        (do ((t l (+ t #e2e5))
             (qs qs (map (lambda (p q) (modulo (- q #e1e5) p)) ps qs)))
            (pi pi)
          (let ((bv (make-vector #e1e5 #t)))
            (do ((qs qs (cdr qs)) (ps ps (cdr ps))) ((null? qs))
              (do ((j (car qs) (+ j (car ps)))) ((<= #e1e5 j))
                (vector-set! bv j #f)))
            (do ((j 0 (+ j 1))) ((= j #e1e5))
              (when (vector-ref bv j) (set! z (+ z 1)))
              (when (= (+ t j j 1) n) (set! pi z)))))))))

(define (nth-prime n)
  (define (bsearch n)
    (let loop ((lo 0) (hi 10000))
      (if (< hi lo) hi
        (let ((mid (quotient (+ lo hi) 2)))
          (cond ((< (vector-ref pi-vector mid) n) (loop (+ mid 1) hi))
                ((< n (vector-ref pi-vector mid)) (loop lo (- mid 1)))
                (else (- mid 1)))))))
  (if (< 37607912018 n) (error 'nth-prime "out of range")
    (if (< n #e1e8) (list-ref (primes #e1e8) (- n 1))
      (let* ((b (bsearch n)) (l (* b #e1e8))
             (ps (cdr (primes (+ (isqrt (+ l #e1e8)) 1))))
             (qs (map (lambda (p) (modulo (* -1/2 (+ l 1 p)) p)) ps))
             (z (vector-ref pi-vector b)) (nth #f))
        (do ((t l (+ t #e2e5))
             (qs qs (map (lambda (p q) (modulo (- q #e1e5) p)) ps qs)))
            (nth nth)
          (let ((bv (make-vector #e1e5 #t)))
            (do ((qs qs (cdr qs)) (ps ps (cdr ps))) ((null? qs))
              (do ((j (car qs) (+ j (car ps)))) ((<= #e1e5 j))
                (vector-set! bv j #f)))
            (do ((j 0 (+ j 1))) ((= j #e1e5))
              (when (vector-ref bv j) (set! z (+ z 1)))
              (when (and (not nth) (= z n)) (set! nth (+ t j j 1))))))))))