(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) ...))))))]))

  )