#lang racket

(require data/queue)
(require racket/generator)
(provide 
 (contract-out
  (make-cmwc-gen (-> (listof integer?) integer? integer? integer? (and/c generator? (-> integer?))))
  (make-default-cmwc-gen (-> integer? (and/c generator? (-> integer?))))
  (make-cmwc-gen-raw (-> queue? integer? integer? integer? (and/c generator? (-> integer?))))
  (init-cmwc-seed (-> integer? integer? queue?))))

(define PHI 2619875739)

(define (cmwc! x a b c)
  (let* ((x0 (dequeue! x))
         (x-new (modulo (- (sub1 b) (+ (* a x0) c)) b))
         (c-new (quotient (+ (* a x0) c) b)))
    (enqueue! x x-new)
    (values x-new c-new)))

(define (cmwc-default x a b c)
  (let* ((x0 (dequeue! x))
         (x-new (modulo (- (sub1 b) (+ (* a x0) c)) b))
         (c-new (quotient (+ (* a x0) c) b)))
    (enqueue! x x-new)
    (values x-new c-new)))

(define (make-cmwc-gen seed a b c)
  (let ((x-reg (make-queue)))
    (for-each (λ (x)
                (enqueue! x-reg x))
              seed)
    (make-cmwc-gen-raw x-reg a b c)))

(define (make-cmwc-gen-raw seed a b c)
  (generator ()
    (let gloop ((current-c c))
      (let-values (((x-new c-new) (cmwc! seed a b current-c)))
        (yield x-new)
        (gloop c-new)))))

(define (init-cmwc-seed init lag)
  (let ((seed (make-queue))
        (q1 (box init))
        (q2 (box (+ init PHI)))
        (q3 (box (+ init PHI PHI)))
        (q4 (box #f)))
    (begin
      (enqueue! seed (unbox q1))
      (enqueue! seed (unbox q2))
      (enqueue! seed (unbox q3)))
    (for ((i (in-range 3 lag)))
      (begin
        (set-box! q4 (bitwise-xor (unbox q2) (unbox q1) PHI i))
        (set-box! q1 (unbox q2))
        (set-box! q2 (unbox q3))
        (set-box! q3 (unbox q4))
        (enqueue! seed (unbox q4))))
    seed))

(define (make-default-cmwc-gen seed) (make-cmwc-gen-raw (init-cmwc-seed seed 4096) 18782 (expt 2 32) 362436))
; TODO: MAKE DEFAULT GENERATOR FASTER (ARITHMETIC SHIFT)
; TODO: DOCUMENT IT PROPERLY
; TODO: MAKE OUTPUT PASS (psuedo-random-generator? x)
; TODO: TEST DEFAULT BETTER

(define ran-max 4294967087)
 
(define (square x) (* x x))
 
(define (expmod base exp m)
  (cond ((= exp 0) 1)
        ((even? exp) (remainder (square (expmod base (/ exp 2) m)) m))
        (else (remainder (* base (expmod base (- exp 1 ) m)) m))))
 
(define (power-of-two-factor x)
  (define (power-of-two-internal x n)
    (if (= (modulo x 2) 1)
        n
        (power-of-two-internal (/ x 2) (+ n 1))))
  (power-of-two-internal x 0))
 
(define (factor-even-number x)
  (let ((po2 (power-of-two-factor x)))
    (values po2 (/ x (expt 2 po2)))))
 
(define (big-random n)
  (let ((r (round (/ (* n (random ran-max)) ran-max))))
    (if (= r 1) (big-random n) r)))
 
(define (num-repetitions p)
  (define 1-p (- 1 p))
  (ceiling (- (/ (log 1-p) (log 4)))))



(define (is-prime? n prob)
  (cond
    ((< n 5) #f)
    ((even? n) #f)
    (else
     (define-values (s d) (factor-even-number (- n 1)))
     (define k (num-repetitions prob))    
     (define (outer-loop k)
       (define a (big-random (- n 2)))
       (define x (expmod a d n))
       (cond
         ((= k 0) #t)
         ((or (= x 1) (= x (- n 1))) (outer-loop (- k 1)))
         (else (inner-loop (- s 1) x))))    
     (define (inner-loop i y)
       (set! y (modulo (* y y) n))
       (cond
         ((or (= y 1) (= i 0)) #f)
         ((= y (- n 1)) (outer-loop (- k 1)))
         (else (inner-loop (- i 1) y))))     
     (outer-loop k))))