random.ss
(module random mzscheme

(require (lib "contract.ss")
(lib "etc.ss")
(lib "plt-match.ss")
(planet "combinators.ss" ("cce" "combinators.plt" 1 4))
(prefix schematics:
(planet "random.ss" ("schematics" "random.plt" 1 0))))

(require-for-syntax
(planet "syntax-utils.ss" ("cce" "syntax-utils.plt" 1 1)))

;; A (Generator T) is (make-generator (-> T)).
(define-struct generator (proc))

;; A (Gen T) is either a T or a (Generator T).
(define (gen/c T) (or/c generator? T))

;; A Nat is a natural number.
(define nat/c natural-number/c)

;; A PosInt is a positive integer.
(define pos-int/c (and/c integer? positive?))

;; A Probability is a real number in the interval [0,1].
(define prob/c (real-in 0 1))

;; A Weight is a positive real number.
(define weight/c (>/c 0))

;; A Fun is any procedure producing a single value.
(define fun/c (unconstrained-domain-> any/c))

;; A GenFun is a function producing a generator.
(define gen-fun/c (unconstrained-domain-> (gen/c any/c)))

;; WeightedGens are a list of weights and generators.
(define weighted-gens/c
(flat-rec-contract weighted-gens/c
null?
(cons/c weight/c (cons/c (gen/c any/c) weighted-gens/c))))

(define (integer-above/c lo)
(flat-named-contract (format "an integer >= ~s" lo)
(lambda (hi) (and (integer? hi) (>= hi lo)))))

(provide/contract

[default-generate-attempts (parameter/c pos-int/c)]

[generator? (-> any/c boolean?)]
[generate (opt-> [(gen/c any/c)] [(-> any/c any/c) pos-int/c] any/c)]
[nonrandom (-> any/c generator?)]

[choose-int-between (->r ([lo integer?]
[hi (integer-above/c lo)])
integer?)]
[random-int-between (->r ([lo integer?]
[hi (integer-above/c lo)])
generator?)]

[choose-size (opt-> [] [nat/c pos-int/c] nat/c)]
[random-size (opt-> [] [nat/c pos-int/c] generator?)]

[choose-boolean (opt-> [] [prob/c] boolean?)]
[random-boolean (opt-> [] [prob/c] generator?)]

[choose-char (opt-> [] [(gen/c nat/c)] char?)]
[random-char (opt-> [] [(gen/c nat/c)] generator?)]

[choose-group-of (opt-> [fun/c (gen/c any/c)] [(gen/c nat/c)] any/c)]
[random-group-of (opt-> [fun/c (gen/c any/c)] [(gen/c nat/c)] generator?)]

[choose-list-of (opt-> [(gen/c any/c)] [(gen/c nat/c)] list?)]
[random-list-of (opt-> [(gen/c any/c)] [(gen/c nat/c)] generator?)]

[choose-vector-of (opt-> [(gen/c any/c)] [(gen/c nat/c)] vector?)]
[random-vector-of (opt-> [(gen/c any/c)] [(gen/c nat/c)] generator?)]

[choose-string (opt-> [] [(gen/c char?) (gen/c nat/c)] string?)]
[random-string (opt-> [] [(gen/c char?) (gen/c nat/c)] generator?)]

[choose-bytes (opt-> [] [(gen/c byte?) (gen/c nat/c)] bytes?)]
[random-bytes (opt-> [] [(gen/c byte?) (gen/c nat/c)] generator?)]

[choose-apply (->* [fun/c] (listof (gen/c any/c)) [any/c])]
[random-apply (->* [fun/c] (listof (gen/c any/c)) [generator?])]

[choose-list (->* [] (listof (gen/c any/c)) [list?])]
[random-list (->* [] (listof (gen/c any/c)) [generator?])]

[choose-vector (->* [] (listof (gen/c any/c)) [vector?])]
[random-vector (->* [] (listof (gen/c any/c)) [generator?])]

[choose-symbol (opt-> [] [(gen/c string?)] symbol?)]
[random-symbol (opt-> [] [(gen/c string?)] generator?)]

[choose-uniform (->* [] (listof (gen/c any/c)) [any/c])]
[random-uniform (->* [] (listof (gen/c any/c)) [generator?])]

[choose-weighted (->* [] weighted-gens/c [any/c])]
[random-weighted (->* [] weighted-gens/c [generator?])]

[choose-weighted* (-> (listof (cons/c weight/c (gen/c any/c))) any/c)]
[random-weighted* (-> (listof (cons/c weight/c (gen/c any/c))) generator?)]

[choose-function (-> gen-fun/c fun/c)]
[random-function (-> gen-fun/c generator?)]
)

(provide
define-generator
let*-random
choose-recursive
random-recursive)

;; make-distribution : (A ... -> B) -> (A ... -> (Generator B))
;; Creates a distribution out of a function.
(define (make-distribution f)
(lambda args (make-generator (lambda () (apply f args)))))

;; default-generate-attempts : (Parameter PosInt)
;; The default number of attempts for generating a random value.
(define default-generate-attempts (make-parameter 100))

;; generate-one : (Gen T) -> T
;; Generates a random value from a distribution.
(define (generate-one gen)
(if (generator? gen)
((generator-proc gen))
gen))

;; generate : (Gen T) [(T -> Boolean) PosInt] -> T
;; Generate values from a generator, until one matches a predicate
;; or the alloted number of tries are done.
(define generate
(opt-lambda (gen [pred (constant #t)]
[max (default-generate-attempts)])
(let loop ([count max]
[tried null])
(if (< count 1)
(error 'generate
"could not find a satisfactory random value; tried:\n~s"
tried)
(let* ([v (generate-one gen)])
(if (pred v) v (loop (- count 1) (cons v tried))))))))

;; nonrandom : T -> (Generator T)
;; Creates a constant distribution.
(define (nonrandom v)
(make-generator (constant v)))

;; choose-int-between : Int Int -> (Generator Int)
;; Choose an integer uniformly distributed between i and j, inclusive.
(define (choose-int-between i j)
(let* ([lo (min i j)]
[hi (max i j)]
[diff (- hi lo)])
(+ (schematics:random-integer (+ diff 1)) lo)))

(define random-int-between (make-distribution choose-int-between))

;; choose-size : [PosInt Nat] -> Nat
;; Choose a natural number with a geometric distribution.
;; NOTE: very inefficient, especially for high average,
;; but much more accurate than schematics:random-geometric.
;; Can probably be made more efficient without sacrificing accuracy.
(define choose-size
(opt-lambda ([minimum 0] [average 4])
(let* ([prob (/ 1 (+ average 1))])
(let loop ([base minimum])
(if (choose-boolean prob)
base
(loop (+ base 1)))))))

(define random-size (make-distribution choose-size))

;; choose-boolean : [Prob] -> Boolean
;; Choose true with the given probability, false otherwise.
(define choose-boolean
(opt-lambda ([prob 1/2])
(if (< (schematics:random-real) prob) #t #f)))

(define random-boolean (make-distribution choose-boolean))

;; choose-char : [(Gen Nat)] -> Char
;; Choose a character with the given distribution of integer code.
(define choose-char
(opt-lambda ([code-gen (random-int-between
(char->integer #\A)
(char->integer #\Z))])
(integer->char (generate code-gen))))

(define random-char (make-distribution choose-char))

;; choose-group-of : (T ... -> (Groupof T)) (Gen T) [(Gen Nat)] -> (Groupof T)
;; Construct a group of a random number of randomly generated elements.
(define choose-group-of
(opt-lambda (make elem-gen [len-gen (random-size)])
(apply make (build-list (generate len-gen)
(lambda (i) (generate elem-gen))))))

(define random-group-of (make-distribution choose-group-of))

;; choose-list-of : (Gen T) [(Gen Nat)] -> (Listof T)
;; Construct a list of a random number of randomly generated elements.
(define choose-list-of
(curry choose-group-of list))

(define random-list-of (make-distribution choose-list-of))

;; choose-vector-of : (Gen T) [(Gen Nat)] -> (Vectorof T)
;; Construct a vector of a random number of randomly generated elements.
(define choose-vector-of
(curry choose-group-of vector))

(define random-vector-of (make-distribution choose-vector-of))

;; choose-string : [(Gen Char) (Gen Nat)] -> String
;; Choose a string with a random number of randomly generated characters.
(define choose-string
(opt-lambda ([char-gen (random-char)] [len-gen (random-size)])
(choose-group-of string char-gen len-gen)))

(define random-string (make-distribution choose-string))

;; choose-bytes : [(Gen Byte) (Gen Nat)] -> Bytes
;; Choose a byte-string with a random number of randomly generated bytes.
(define choose-bytes
(opt-lambda ([byte-gen (random-int-between 0 255)] [len-gen (random-size)])
(choose-group-of bytes byte-gen len-gen)))

(define random-bytes (make-distribution choose-bytes))

;; choose-apply : (A ... -> B) (Gen A) ... -> (Gen B)
;; Apply a function to randomly generated inputs.
(define (choose-apply f . gens)
(apply f (map generate gens)))

(define random-apply (make-distribution choose-apply))

;; choose-list : (Gen T) ... -> (Generator (list T ...))
;; Construct a list with a fixed number of randomly generated elements.
(define (choose-list . gens)
(map generate gens))

(define random-list (make-distribution choose-list))

;; choose-vector : (Gen T) ... -> (Generator (vector T ...))
;; Construct a vector with a fixed number of randomly generated elements.
(define (choose-vector . gens)
(apply vector (map generate gens)))

(define random-vector (make-distribution choose-vector))

;; choose-symbol : [(Gen String)] -> Symbol
;; Choose a symbol with a randomly generated name.
(define choose-symbol
(opt-lambda ([name-gen (random-string (random-char) (random-size 1))])
(string->symbol (generate name-gen))))

(define random-symbol (make-distribution choose-symbol))

;; choose-uniform : (Gen T1) ... (Gen TK) -> (Or T1 ... TK)
;; Choose a value from a fairly-chosen random distribution.
(define (choose-uniform . gens)
(let* ([v (apply vector gens)]
[c (vector-length v)])
(generate (vector-ref v (schematics:random-integer c)))))

(define random-uniform (make-distribution choose-uniform))

;; cons-weights-and-gens : WeightedGens -> (Listof (cons Weight Gen))
;; Pairs weights and generators.
(define (cons-weights-and-gens alternation)
(match alternation
[(list) (list)]
[(list-rest weight gen rest)
(cons (cons weight gen) (cons-weights-and-gens rest))]))

;; choose-weighted : Weight1 (Gen T1) ... WeightK (Gen TK) -> (Or T1 ... TK)
;; A convenient alternate calling convention for choose-weighted*.
(define (choose-weighted . weights-and-gens)
(choose-weighted* (cons-weights-and-gens weights-and-gens)))

(define random-weighted (make-distribution choose-weighted))

;; choose-weighted* : (list (cons Weight (Gen T)) ...) -> (Or T ...)
;; Choose a value from one of the random distributions,
;; chosen with the corresponding weight.
(define (choose-weighted* pairs)
(let* ([total (apply + (map car pairs))]
[choice (schematics:random-real)])
(let loop ([base 0]
[pairs pairs])
(match pairs
[(list (cons weight gen)) (generate gen)]
[(list-rest (cons weight gen) rest)
(let* ([sum (+ base weight)]
[prob (/ sum total)])
(if (<= choice prob)
(generate gen)
(loop sum rest)))]))))

(define random-weighted* (make-distribution choose-weighted*))

;; choose-function : (A ... -> (Gen B)) -> (A ... -> B)
;; Constructs a function which generates a random output
;; for each unique set of inputs.
(define (choose-function f)
(let* ([table (make-hash-table 'equal)])
(lambda args
(hash-table-get
table args
(lambda ()
(let* ([result (generate (apply f args))])
(hash-table-put! table args result)
result))))))

(define random-function (make-distribution choose-function))

;; (random-recursive name [weight gen] ...) : Macro
;; Constructs a recursively-defined distribution with weighted alternatives.
(define-syntax (random-recursive stx)
(syntax-case stx ()
[(r-r name [weight gen] ...)
(syntax/loc stx
(let* ([pairs null]
[name (make-generator (lambda () (choose-weighted* pairs)))])
(set! pairs (list (cons weight gen) ...))
name))]))

;; (choose-recursive name [weight gen] ...) : Macro
;; Chooses from a recursively-defined distribution, as above.
(define-syntax (choose-recursive stx)
(syntax-case stx ()
[(c-r name [weight gen] ...)
(syntax/loc stx
(generate (random-recursive name [weight gen] ...)))]))

;; (let*-random ((var gen [pred count])) body ...) : Macro where [] = optional
;; Binds each var to a random value in body.
(define-syntax (let*-random stx)
(syntax-case stx ()
[(lr ([var gen] . rest) . body)
(syntax/loc stx (lr ([var gen #t] . rest) . body))]
[(lr ([var gen pred] . rest) . body)
(syntax/loc stx
(lr ([var gen pred (default-generate-attempts)] . rest) . body))]
[(lr ([var gen pred count] . rest) . body)
(syntax/loc stx
(let* ([var (generate gen (lambda (var) pred) count)])
(lr rest . body)))]
[(lr () . body)
(syntax/loc stx (let* () . body))]))

(define-syntax (define-generator stx)
(syntax-case stx ()
[(dg (name arg ...) (weight gen) ...)
(syntax/loc stx
(define (name arg ...)
(make-generator
(lambda ()
(choose-weighted* (list (cons weight gen) ...))))))]))

)