;;
;; A port of the SML/NJ implementation of a random number generator
;; using the subtract-with-borrow (SWB) method described by Marsaglia
;; and Zaman in _A New Class of Random Number Generators_,
;; Ann. Applied Prob. 1(3), 1991, pp. 462-480.
;; 
;; The SWB generator is a 30-bit generator with lags 48 and 8. It has
;; period (2^1487 - 2^247)/105 or about 10^445. The SWB generator acts
;; locally like a lagged Fibonacci generator and thus it is combined
;; with a linear congruential generator (48271*a)mod(2^30-1).
;;
;;
;; Copyright 2007-2010 Ivan Raikov.
;;
;; This program is free software: you can redistribute it and/or
;; modify it under the terms of the GNU General Public License as
;; published by the Free Software Foundation, either version 3 of the
;; License, or (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful, but
;; WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;; General Public License for more details.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>."))))


(module random-swb


 (make-swb-random-state
  swb:random!
  swb:random-natural!
  swb:random-real!
  swb:random-range!
  swb:random
  swb:random-natural
  swb:random-real
  swb:random-range)

  (import scheme chicken extras srfi-4 matchable foreign)

  (require-extension srfi-4 matchable)


(define nbits 30)
(define max-word (inexact->exact (- (expt 2 nbits) 1)))

(define bit29 #x20000000)
(define lo29  #x1FFFFFFF)


(define N 48)
(define lag 8)
(define offset (- N lag))


#>
  unsigned int cminus (unsigned int x, unsigned int y, int b)
  {
     if (b > 0) 
     {
      return x-y-1;
      }
     else
     {
      return x-y;
      }
  }
<#

(define cminus (foreign-lambda unsigned-int "cminus" unsigned-int unsigned-int unsigned-int))



(define (swb:error x . rest)
  (let ((port (open-output-string)))
    (let loop ((objs (cons x rest)))
      (if (null? objs)
	  (begin
	    (newline port)
	    (error 'random-swb (get-output-string port)))
	  (begin (display (car objs) port)
		 (display " " port)
		 (loop (cdr objs)))))))

(define 2^-29 (expt 2 -29))

;;
;; Random number generator state representation
;;
;; * seeds:    seed vector
;; * borrow:   last borrow
;; * congseed: congruential seed 
;; * index:    index of next available value in seeds 
;;
(define-record swb-state seeds borrow congseed index)

(define-record-printer (swb-state x out)
  (fprintf out "#(SWB seeds=~S borrow=~S congseed=0x~X index=~S)"
	   (swb-state-seeds x)
	   (swb-state-borrow x)
	   (swb-state-congseed x)
	   (swb-state-index x)))



;;
;; Linear congruential generator:
;; multiplication by 48271 mod (2^30 - 1) 
;;
(define  lcg-a 48271)
(define  lcg-m max-word)
(define  lcg-q (quotient lcg-m lcg-a))
(define  lcg-r (modulo lcg-m lcg-a))

(define (lcg s)
  (let ((left  (fx* lcg-a (modulo s lcg-q)))
	(right (fx* lcg-r (quotient s lcg-q))))
    (if (fx> left right) (fx- left right)
	(fx+ (fx- lcg-m right) left))))

;;
;; Fills seed array using subtract-with-borrow generator:
;;
;;  x[n] = x[n-lag] - x[n-N] - borrow
;; Sets index to 1 and returns 0th value.
;;
(define (fill! state)
  
  (define seeds (swb-state-seeds state))
  (define index (swb-state-index state))
  (define borrow (swb-state-borrow state))

  (define (update ix iy b)
    (let ((x (u32vector-ref seeds ix))
	  (y (u32vector-ref seeds iy)))
      (let ((z (cminus x y (if b 1 0))))
	(u32vector-set! seeds iy z)
	(if b (fx>= y x) (fx> y x)))))

  (define (fillup i b)
    (if (fx= i lag) b
	(fillup (fx+ 1 i) (update (fx+ i offset) i b))))

  (define (fillup1 i b)
    (if (fx= i N) b
	(fillup1 (fx+ 1 i) (update (fx- i lag) i b))))

  (swb-state-borrow-set! state (fillup1 lag (fillup 0 borrow)))
  (swb-state-index-set! state 1)
  (u32vector-ref seeds 0))

;;
;; Creates an initial seed vector and generator state.  Fills the seed
;; array one bit at a time by taking the leading bit of the xor of a
;; shift register and a congruential sequence.  The congruential
;; generator is (c*48271) mod (2^30 - 1).  The shift register
;; generator is c(I + L18)(I + R13).  The same congruential generator
;; continues to be used as a mixing generator with the SWB generator.
;;
(define (make-swb-random-state congy shrgx)

  (define (mki v)
    (match v ((i c s)
	      (let* ((c1 (lcg  c)))
		(let* ((s1 (bitwise-xor s  (fxshl s  18)))
		       (s2 (bitwise-xor s1 (fxshr s1 13)))
		       (i1 (fx+ (bitwise-and lo29 (fxshr i 1))
				(bitwise-and bit29 (bitwise-xor c1 s2)))))
		  (list i1 c1 s2))))
	      (else (swb:error 'mki "invalid argument " v))))

  (define (iterate n v)
    (if (zero? n) v
	(iterate (fx- n 1) (mki v))))

  (define (mkseed congx shrgx)
    (iterate nbits (list 0 congx shrgx)))

  (define (genseed n seeds congx shrgx)
    (if (zero? n) (values seeds congx)
	(let ((v (mkseed congx shrgx)))
	  (match v  ((seed congx1 shrgx1)
		     (genseed (fx- n 1) (cons seed seeds) congx1  shrgx1))
		 (else (swb:error 'genseed "invalid mkseed return value " v))))))

  (define congx (fx+ 1 (fxshl (bitwise-and congy max-word) 1)))

  (let-values (((seeds congx)  (genseed N (list) congx shrgx)))
    (make-swb-state (list->u32vector seeds) #f congx 0)))

;;
;; Computes the next random number. The tweak function xor's the
;; number from the SWB generator with a number from the linear
;; congruential generator.
;;
(define (swb:random! s)
  (match s 
	 (($ swb-state seeds borrow congseed index)
	  (let ((tweak (lambda (i) 
			 (let ((c (lcg congseed)))
			   (swb-state-congseed-set! s c)
			   (bitwise-xor i c)))))
	    (if (fx= index N)
		(tweak (fill! s))
		(begin
		  (swb-state-index-set! s (fx+ 1 index))
		  (tweak (u32vector-ref seeds index))))))
	 (else (swb:error 'swb:random "invalid state " s))))

(define (swb:random-natural! s)
  (bitwise-and (swb:random! s) lo29))

(define (swb:random-real! s)
  (let* ((n1 (swb:random-natural! s))
	 (n2 (swb:random-natural! s)))
    (* 2^-29 (+ n1 (* n2 2^-29)))))

(define (swb:random-range! i j)
  (if (fx< j i)
      (swb:error 'swb:random-range "argument j < i")
      (if (and (fixnum? i) (fixnum? j))
	  (let ((R  (* 2^-29 (fx+ 1 (fx- j i)))))
	    (lambda (s) (fx+ i (inexact->exact (truncate (* R (swb:random-natural! s)))))))
	  (let ((R  (fx+ 1 (fx- j i))))
	    (lambda (s) (+ i (* R (swb:random-real! s))))))))
	  
	  
;;
;; Side-effect-free variants of the above functions.
;;
(define (swb:random s)
  (match s 
	 (($ swb-state seeds borrow congseed index)
	  (let ((tweak (lambda (s i) 
			 (let ((c (lcg congseed)))
			   (swb-state-congseed-set! s c)
			   (bitwise-xor i c)))))
	    (if (fx= index N)
		(let* ((s1  (make-swb-state seeds borrow congseed index))
		       (i   (fill! s1))
		       (v   (tweak s1 i)))
		  (values s1 v))
		(let* ((s1 (make-swb-state seeds borrow congseed (fx+ 1 index)))
		       (v  (tweak s1 (u32vector-ref seeds index))))
		  (values s1 v)))))
	 (else (swb:error 'swb:random "invalid state " s))))

(define (swb:random-natural s)
  (let-values (((s v) (swb:random s)))
    (values s (bitwise-and v lo29))))

(define (swb:random-real s)
  (let-values (((s1 n1) (swb:random-natural s)))
    (let-values (((s2 n2) (swb:random-natural s1)))
       (values s2 (* 2^-29 (+ n1 (* n2 2^-29)))))))

(define (swb:random-range i j)
  (if (fx< j i)
      (swb:error 'swb:random-range "argument j < i")
      (if (and (fixnum? i) (fixnum? j))
	  (let ((R  (* 2^-29 (fx+ 1 (fx- j i)))))
	    (lambda (s) 
	      (let-values (((s1 v) (swb:random-natural s)))
		(values s1 (fx+ i (inexact->exact (truncate (* R v))))))))
	  (let ((R  (fx+ 1 (fx- j i))))
	    (lambda (s) 
	      (let-values (((s1 v)  (swb:random-real s)))
	          (values s1 (+ i (* R v)))))))))
	  
	  
)