;;; -*- Mode: LISP; Syntax: ANSI-Common-Lisp; Package: RANDOM -*-
;; A translation of the random module in the Python standard library
;; Translated by Paul Foley <mycroft@actrix.gen.nz>
;; http://public.xdi.org/=pf

#+CMU (ext:file-comment "$Header$")

(defpackage "RANDOM"
  (:use "COMMON-LISP")
  (:export "RANDOM" "UNIFORM" "NORMALVARIATE" "LOGNORMALVARIATE"
	   "CUNIFVARIATE" "EXPOVARIATE" "VONMISESVARIATE" "GAMMAVARIATE"
	   "BETAVARIATE" "PARETOVARIATE" "WEIBULLVARIATE"
	   "CHOICE" "SHUFFLE" "SAMPLE"))

(provide :random)

(in-package "RANDOM")

(defconstant e 2.71828182845904523536028747135266249775724709369996l0)
(defconstant +NV-MAGIC-CONSTANT+ (* 4 (/ (exp -0.5d0) (sqrt 2d0))))
(defconstant +SG-MAGIC-CONSTANT+ (1+ (log 4.5d0)))

(declaim (ftype (function () (double-float 0d0 (1d0))) random1))
(defun random1 ()
  (random 1d0))

(defun uniform (a b)
  (+ a (* (- b a) (random1))))

(defun normalvariate (mu sigma)
  (loop for u1 = (random1)
	for u2 = (- 1d0 (random1))
	for z = (/ (* +NV-MAGIC-CONSTANT+ (- u1 0.5d0)) u2)
	for zz = (/ (* z z) 4d0)
    until (<= zz (- (log u2)))
    finally (return (+ mu (* z sigma)))))

(defun lognormalvariate (mu sigma)
  (exp (normalvariate mu sigma)))

(defun cunifvariate (mu arc)
  (mod (+ mu (* arc (- (random1) 0.5d0))) pi))

(defun expovariate (lambda)
  (loop for u = (random1) while (<= u 1d-7)
    finally (return (- (/ (log u) lambda)))))

(defun vonmisesvariate (mu kappa)
  (loop with a = (1+ (sqrt (1+ (* 4d0 kappa kappa))))
	with b = (- a (/ (sqrt (* 2d0 a)) (* 2d0 kappa)))
	with r = (/ (1+ (* b b)) (* 2d0 b))
	for u1 = (random1)
	for z = (cos (* pi u1))
	for f = (/ (1+ (* r z)) (+ r z))
	for c = (* kappa (- r f))
	for u2 = (random1)
	until (and (>= u2 (* c (- 2d0 c)))
		   (> u2 (* c (exp (- 1d0 c)))))
   finally (let ((u3 (random1)))
	     (if (> u3 0.5d0)
		 (return (+ (mod mu (* 2 pi)) (acos f)))
		 (return (- (mod mu (* 2 pi)) (acos f)))))))

(defun gammavariate (alpha beta)
  (cond ((> alpha 1d0)
	 (loop with ainv = (sqrt (- (* 2d0 alpha) 1d0))
	       with bbb = (- alpha (log 4d0))
	       with ccc = (+ alpha ainv)
	       for u1 = (loop for u1 = (random1) until (< 1d-7 u1 0.9999999)
			  finally (return u1))
	       for u2 = (- 1d0 (random1))
	       for v = (/ (log (/ u1 (- 1d0 u1))) ainv)
	       for x = (* alpha (exp v))
	       for z = (* u1 u1 u2)
	       for r = (+ bbb (* ccc v) (- x))
	   until (or (>= (- (+ r +SG-MAGIC-CONSTANT+) (* 4.5d0 z)) 0d0)
		     (>= r (log z)))
	   finally (return (* x beta))))
	((= alpha 1d0)
	 (let ((u (loop for u = (random1) while (<= u 1d-7)
		    finally (return u))))
	   (* (- (log u)) beta)))
	(t
	 (loop for u = (random1)
	       for b = (/ (+ e alpha) e)
	       for p = (* b u)
	       for x = (if (<= p 1d0)
			   (expt p (/ 1d0 alpha))
			   (- (log (/ (- b p) alpha))))
	       for u1 = (random1)
	   until (or (and (<= p 1d0) (> u1 (exp (- x))))
		     (and (> p 1d0)  (> u1 (expt x (- alpha 1d0)))))
	   finally (return (* x beta))))))

;; GAUSS -- a faster version of normalvariate
;;     The Python version uses both cos and sin, returning one value and
;;     caching the other to return on the next call (on the third call it
;;     calculates a new random number)
;;     I'd probably rather return two values, though...
;;(defun gauss (mu sigma)
;;  (let ((a (* 2 pi (random1)))
;;	(b (* sigma (sqrt (* -2 (log (- 1 (random1))))))))
;;    (values (+ mu (* b (cos a)))
;;	    (+ mu (* b (sin a))))))

(defun betavariate (alpha beta)
  (let ((y (gammavariate alpha 1)))
    (if (zerop y)
	0d0
	(/ y (+ y (gammavariate beta 1))))))

(defun paretovariate (alpha)
  (/ 1d0 (expt (- 1d0 (random1)) (/ 1d0 alpha))))

(defun weibullvariate (alpha beta)
  (* alpha (expt (- (log (- 1d0 (random1)))) (/ 1d0 beta))))

(defun choice (sequence)
  (elt sequence (random (length sequence))))

(defun shuffle (sequence)
  (let ((temp (coerce sequence 'vector)))
    (loop for i downfrom (1- (length temp)) to 1 do
      (rotatef (aref temp i) (aref temp (random (1+ i)))))
    (unless (eq temp sequence)
      (replace sequence temp))
    sequence))

(defun sample (population k)
  (let* ((n (length population))
	 (type (if (vectorp population) (array-element-type population) 't))
	 (result (make-array k :element-type type)))
    (if (< n (* 6 k))
	(let ((pool (make-array n :element-type type
				  :initial-contents population)))
	  (loop for i below k as j = (random (- n i)) do
	    (shiftf (aref result i) (aref pool j) (aref pool (- n i 1)))))
	(let ((selected (make-hash-table)))
	  (loop for i below k as j = (random n) do
	    (loop while (gethash j selected) do (setq j (random n)))
	    (setf (aref result i) (elt population j)
		  (gethash j selected) t))))
    result))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defun range (stop/start &optional (stop 0 stop-p) (step 1))
  (declare (type real stop/start stop step))
  (let ((start (truncate (if stop-p stop/start 0)))
	(stop (truncate (if stop-p stop stop/start)))
	(step (truncate step)))
    (cond ((> step 0)
	   (loop for i from start below stop by step collecting i))
	  ((< step 0)
	   (loop for i from start above stop by (abs step) collecting i))
	  (t
	   (error "Step size is zero.")))))

(define-compiler-macro choice (&whole form sequence)
  (cond ((atom sequence)
	 form)
	((eq (first sequence) 'range)
	 (cond ((fourth sequence)
		(if (and (constantp (second sequence))
			 (constantp (third sequence))
			 (constantp (fourth sequence)))
		    `(+ ,(second sequence)
		        (* ,(fourth sequence)
			   (random ,(ceiling (- (third sequence)
						(second sequence))
					     (fourth sequence)))))
		    `(let ((#1=#:start ,(second sequence))
			   (#2=#:step ,(fourth sequence)))
		       (+ #1# (* #2# (random (ceiling (- ,(third sequence) #1#)
						      #2#)))))))
	       ((third sequence)
		(if (and (constantp (second sequence))
			 (constantp (third sequence)))
		    `(+ ,(second sequence) (random ,(- (third sequence)
						       (second sequence))))
		    `(let ((#3=#:start ,(second sequence)))
		       (+ #3# (random (- ,(third sequence) #3#))))))
	       (t
		`(random ,(second sequence)))))
	;; other optimizations here
	(t form)))

(defun %sample-range (start stop k)
  (let* ((n (- stop start))
	 (result (make-array k)))
    (if (< n (* 6 k))
	(let ((pool (make-array n)))
	  (dotimes (i n)
	    (setf (aref pool i) (+ i start)))
	  (loop for i below k as j = (random (- n i)) do
	    (shiftf (aref result i) (aref pool j) (aref pool (- n i 1)))))
	(let ((selected (make-hash-table)))
	  (loop for i below k as j = (random n) do
	    (loop while (gethash j selected) do (setq j (random n)))
	    (setf (aref result i) (+ start j)
		  (gethash j selected) t))))
    result))

(defun %sample-range/step (start stop step k)
  (let* ((n (ceiling (- stop start) step))
	 (result (make-array k)))
    (if (< n (* 6 k))
	(let ((pool (make-array n)))
	  (dotimes (i n)
	    (setf (aref pool i) (+ (* i step) start)))
	  (loop for i below k as j = (random (- n i)) do
	    (shiftf (aref result i) (aref pool j) (aref pool (- n i 1)))))
	(let ((selected (make-hash-table)))
	  (loop for i below k as j = (random n) do
	    (loop while (gethash j selected) do (setq j (random n)))
	    (setf (aref result i) (+ start (* step j))
		  (gethash j selected) t))))
    result))

(define-compiler-macro sample (&whole form population k)
  (cond ((atom population)
	 form)
	((eq (first population) 'range)
	 (cond ((fourth population)
		`(%sample-range/step ,(second population) ,(third population)
				     ,(fourth population) ,k))
	       ((third population)
		`(%sample-range ,(second population) ,(third population) ,k))
	       (t
		`(%sample-range 0 ,(second population) ,k))))
	;; other optimizations here
	(t form)))