(import chicken.random)  ; for pseudo-random-integer

(define (square x) (* x x))

(define (non-trivial-root? x n)
  (and (not (= x (- n 1)))
       (not (= x 1))
       (= (remainder (square x) n) 1)))

(define (expmod base exp m)
  (cond ((= exp 0) 1)
	((even? exp)
	  (let ((test (expmod base (/ exp 2) m)))
	    (if (non-trivial-root? test m)
		0  ; Signal non-trivial root
	        (remainder (square test)
		           m))))
	 (else
	  (remainder (* (expmod base (- exp 1) m) base)
		     m))))

(define (miller-rabin-test n count)
  (let* ((a (+ 1 (pseudo-random-integer (- n 1))))
	 (result (expmod a (- n 1) n)))
    ; (display a)
    ; (newline)
    (cond ((= count 0) 
	      ; (display "Count 0")
	      ; (newline)
   	      #t)	; Exhausted our attempts.  Assumed prime
	  ((= result 0) 
	      ; (display "Non-trivial root")
	      ; (newline)
	      #f)	; Non-trivial root: not prime
	  ((> result 1) 
	      ; (display "MR test failed: ")
	      ; (display result)
	      ; (newline)
	      #f)	; a^(n-1) not congruent 1 (n): not prime
	  (else (miller-rabin-test n (- count 1))))))


(define (prime? n)
  (miller-rabin-test n 25))