#lang racket

;;; dds/utils

;;; Various utilities.

(require
 graph
 (for-syntax syntax/parse racket/list))

(provide
 ;; Functions
 (contract-out [unstringify-pairs (-> (listof (general-pair/c string? any/c))
                                      (listof (general-pair/c symbol? any/c)))]
               [read-org-variable-mapping (-> string? variable-mapping?)]
               [unorgv (-> string? variable-mapping?)]
               [dotit (-> graph? void?)]
               [read-symbol-list (-> string? (listof symbol?))]
               [drop-first-last (-> string? string?)]
               [list-sets->list-strings (-> (listof (set/c any/c)) (listof string?))]
               [pretty-print-set-sets (-> (set/c (set/c symbol?) #:kind 'dont-care) string?)]
               [update-vertices/unweighted (-> graph? (-> any/c any/c) graph?)]
               [update-graph (->* (graph?)
                                  (#:v-func (-> any/c any/c)
                                   #:e-func (-> any/c any/c))
                                  graph?)]
               [pretty-print-set (-> generic-set? string?)]
               [collect-by-key (-> (listof any/c) (listof any/c) (values (listof any/c) (listof (listof any/c))))]
               [collect-by-key/sets (-> (listof any/c) (listof any/c) (values (listof any/c) (listof (set/c any/c))))]

               [ht-values/list->set (-> (hash/c any/c (listof any/c)) (hash/c any/c (set/c any/c)))]
               [hash->list/ordered (-> hash? (listof (cons/c any/c any/c)))]
               [multi-split-at (-> (listof (listof any/c)) number?
                                   (values (listof (listof any/c)) (listof (listof any/c))))]
               [lists-transpose (-> (listof (listof any/c)) (listof (listof any/c)))]
               [procedure-fixed-arity? (-> procedure? boolean?)]
               [in-random (case-> (-> (stream/c (and/c real? inexact? (>/c 0) (</c 1))))
                                  (-> (integer-in 1 4294967087) (stream/c exact-nonnegative-integer?))
                                  (-> exact-integer? (integer-in 1 4294967087)
                                      (stream/c exact-nonnegative-integer?)))]
               [cartesian-product/stream (->* () #:rest (listof stream?) stream?)]
               [boolean-power (-> number? (listof (listof boolean?)))]
               [boolean-power/stream (-> number? (stream/c (listof boolean?)))]
               [any->01 (-> any/c (or/c 0 1))]
               [01->boolean (-> (or/c 0 1) boolean?)]
               [hash-intersect (->* [(and/c hash? immutable?)]
                                    [#:combine (-> any/c any/c any/c)
                                     #:combine/key (-> any/c any/c any/c any/c)]
                                    #:rest (listof hash?)
                                    (and/c hash? immutable?))])
 ;; Contracts
 (contract-out [variable-mapping? contract?]
               [string-variable-mapping? contract?]
               [general-pair/c (-> contract? contract? contract?)]))

(module+ test
  (require rackunit))


;;; ===================
;;; HashTable Injection
;;; ===================

;;; This section of the file contains some utilities to streamline the
;;; usage of hash tables mapping symbols to values.  The goal is
;;; essentially to avoid having to write explicit hash-ref calls.

;;; A variable mapping is a hash table mapping symbols to values.
(define (variable-mapping? dict) (hash/c symbol? any/c))

;;; Given a (HashTable Symbol a) and a sequence of symbols, binds
;;; these symbols to the values they are associated to in the hash
;;; table, then puts the body in the context of these bindings.
;;;
;;; > (let ([ht #hash((a . 1) (b . 2))])
;;;     (auto-hash-ref/explicit (ht a b) (+ a (* 2 b))))
;;; 5
;;;
;;; Note that only one expression can be supplied in the body.
(define-syntax (auto-hash-ref/explicit stx)
  (syntax-parse stx
    [(_ (ht:id xs:id ...) body:expr)
     #`(let #,(for/list ([x (syntax->list #'(xs ...))])
                #`[#,x (hash-ref ht '#,x)])
         body)]))

(module+ test
  (test-case "auto-hash-ref/explicit"
    (define mytable #hash((a . 3) (b . 4)))
    (check-equal? (auto-hash-ref/explicit (mytable b a)
                                          (* a b))
                  12)
    (define ht #hash((a . #t) (b . #f)))
    (check-equal? (auto-hash-ref/explicit (ht a b)
                                          (and (not a) b))
                  #f)))

;;; Given an expression and a (HashTable Symbol a), looks up the
;;; symbols with a leading semicolon and binds them to the value they
;;; are associated to in the hash table.
;;;
;;; > (let ([ht #hash((a . 1) (b . 2))])
;;;      (auto-hash-ref/: ht (+ :a (* 2 :b))))
;;; 5
;;;
;;; Note that the symbol :a is matched to the key 'a in the hash
;;; table.
;;;
;;; Note that only one expression can be supplied in the body.
(define-syntax (auto-hash-ref/: stx)
  (syntax-parse stx
    [(_ ht:id body)
     (let* ([names/: (collect-colons (syntax->datum #'body))])
       #`(let #,(for/list ([x names/:])
                  ;; put x in the same context as body
                  #`[#,(datum->syntax #'body x)
                     (hash-ref ht '#,(strip-colon x))])
           body))]))

(module+ test
  (test-case "auto-hash-ref/:"
    (define ht1 #hash((x . #t) (y . #t) (t . #f)))
    (define z #t)
    (check-equal? (auto-hash-ref/: ht1
                                   (and :x (not :y) z (or (and :t) :x)))
                  #f)
    (define ht2 #hash((a . 1) (b . 2)))
    (check-equal? (auto-hash-ref/: ht2 (+ :a (* 2 :b)))
                  5)))

;;; The helper functions for auto-hash-ref/:.
(begin-for-syntax
  ;; Collect all the symbols starting with a colon in datum.
  (define (collect-colons datum)
    (remove-duplicates
     (flatten
      (for/list ([token datum])
        (cond
          [(symbol? token)
           (let ([name (symbol->string token)])
             (if (eq? #\: (string-ref name 0))
                 token
                 '()))]
          [(list? token)
           (collect-colons token)]
          [else '()])))))

  ;; Strip the leading colon off x.
  (define (strip-colon x)
    (let ([x-str (symbol->string x)])
      (if (eq? #\: (string-ref x-str 0))
          (string->symbol (substring x-str 1))
          x))))

;;; Temporarily injects the mappings from the given hash table as
;;; bindings in a namespace including racket/base and then evaluates
;;; the expression.
;;;
;;; > (let ([ht #hash((a . 1) (b . 1))])
;;;     (eval-with ht '(+ b a 1)))
;;; 3
;;;
;;; The local bindings from the current lexical scope are not
;;; conserved.  Therefore, the following outputs an error about a
;;; missing identifier:
;;;
;;; > (let ([ht #hash((a . 1) (b . 1))]
;;;         [z 1])
;;;     (eval-with ht '(+ b z a 1)))
;;;
(define (eval-with ht expr)
  (parameterize ([current-namespace (make-base-namespace)])
    (for ([(x val) ht]) (namespace-set-variable-value! x val))
    (eval expr)))

(module+ test
  (test-case "eval-with"
    (check-equal? (let ([ht #hash((a . 1) (b . 1))])
                    (eval-with ht '(+ b a 1)))
                  3)))

;;; Same as eval-with, but returns only the first value produced by
;;; the evaluated expression.
(define (eval-with1 ht expr)
  (let ([vals (call-with-values (λ () (eval-with ht expr))
                                (λ vals vals))])
    (car vals)))


;;; ==============================
;;; Analysis of quoted expressions
;;; ==============================

;;; Produces a list of symbols appearing in the quoted expression
;;; passed in the first argument.
(define (extract-symbols form)
  (match form
    [(? symbol?) (list form)]
    [(? list?) (flatten (for/list ([x form])
                          (extract-symbols x)))]
    [else '()]))

(module+ test
  (test-case "extract-symbols"
    (check-equal? (extract-symbols '(1 (2 3) x (y z 3)))
                  '(x y z))))


;;; =========================
;;; Interaction with Org-mode
;;; =========================

;;; Org-mode supports laying out the output of code blocks as tables,
;;; which is very practical for various variable mappings (e.g.,
;;; states).  However, when the hash table maps variables to lists,
;;; Org-mode will create a column per list element, which may or may
;;; not be the desired effect.  In this section I define some
;;; utilities for nicer interoperation with Org-mode tables.  I also
;;; define some shortcuts to reduce the number of words to type when
;;; using dds with Org-mode.  See example.org for examples of usage.

;;; Converts any value to string.
(define (any->string x)
  (with-output-to-string (λ () (display x))))

(module+ test
  (test-case "any->string"
   (check-equal? (any->string 'a) "a")
   (check-equal? (any->string '(a 1 (x y))) "(a 1 (x y))")
   (check-equal? (any->string "hello") "hello")))

;;; A string variable mapping is a mapping from variables to strings.
(define (string-variable-mapping? dict) (hash/c symbol? string?))

;;; Converts all the values of a variable mapping to string.
(define (stringify-variable-mapping ht)
  (for/hash ([(key val) ht]) (values key (any->string val))))

(module+ test
  (test-case "stringify-variable-mapping"
    (define mp (stringify-variable-mapping #hash((a . (and a b)) (b . (not b)))))
    (check-equal? (hash-ref mp 'a) "(and a b)")
    (check-equal? (hash-ref mp 'b) "(not b)")))

;;; Reads any value from string.
(define (string->any str)
  (with-input-from-string str (λ () (read))))

(module+ test
  (test-case "string->any"
    (check-equal? (string->any "(or b (not a))") '(or b (not a)))
    (check-equal? (string->any "14") 14)))

;;; Given a sexp, converts all "#f" to #f and "#t" to #t.
;;;
;;; When I read Org-mode tables, I pump them through a call to the
;;; prin1 because the elisp sexp seems incompatible with Racket.  On
;;; the other hand, Racket Booleans seem to upset elisp a little, so
;;; prin1 wraps them in additional double quotes.  This function
;;; removes those quotes.
(define/match (handle-org-booleans datum)
  [("#t") #t]
  [("#f") #f]
  [((? list?)) (map handle-org-booleans datum)]
  [ (_) datum])

;;; Given a sexp, applies the given function to any object which is
;;; not a list.
;;;
;;; The contract of this function will not check whether func is
;;; indeed applicable to every non-list element of the sexp.  If this
;;; is not the case, a contract violation for func will be generated.
(define (map-sexp func sexp)
  (match sexp
    [(? list?) (map ((curry map-sexp) func) sexp)]
    [datum (func datum)]))

(module+ test
  (test-case "map-sexp"
    (check-equal? (map-sexp add1 '(1 2 (4 10) 3)) '(2 3 (5 11) 4))))

;;; Reads a sexp from a string produced by Org-mode for a named table.
;;; See example.org for examples.
(define read-org-sexp
  (compose ((curry map-sexp) (match-lambda
                               [(and (? string?) str) (string->any str)]
                               [x x]))
           string->any))

;;; A shortcut for read-org-sexp.
(define unorg read-org-sexp)

(module+ test
  (test-case "read-org-sexp"
    (check-equal? (read-org-sexp "((\"a\" \"(and a b)\") (\"b\" \"(or b (not a))\"))")
                  '((a (and a b)) (b (or b (not a)))))
    (check-equal? (read-org-sexp "(#t \"#t\" \"#t \" '(1 2 \"#f\"))")
                  '(#t #t #t '(1 2 #f)))))

;;; A contract allowing pairs constructed via cons or via list.
(define (general-pair/c key-contract val-contract)
  (or/c (list/c key-contract val-contract)
        (cons/c key-contract val-contract)))

;;; Given a list of pairs of strings and some other values (possibly
;;; strings), converts the first element of each pair to a string, and
;;; reads the second element with string->any or keeps it as is if it
;;; is not a string.
(define (unstringify-pairs pairs)
  (for/list ([p pairs])
    (match p
      [(list key val)
       (cons (string->symbol key) (if (string? val)
                                      (string->any val)
                                      val))]
      [(cons key val)                   ; also handle improper pairs
       (cons (string->symbol key) (if (string? val)
                                      (string->any val)
                                      val))])))

(module+ test
  (test-case "unstringify-pairs"
    (check-equal? (unstringify-pairs '(("a" . "1") ("b" . "(and a (not b))")))
                  '((a . 1) (b . (and a (not b)))))
    (check-equal? (unstringify-pairs '(("a" . 1) ("b" . "(and a (not b))")))
                  '((a . 1) (b . (and a (not b)))))))

;;; Reads a variable mapping from a string, such as the one which
;;; Org-mode produces from tables.
(define read-org-variable-mapping
  (compose make-immutable-hash unstringify-pairs string->any))

(module+ test
  (test-case "read-org-variable-mapping"
    (define m1 (read-org-variable-mapping "((\"a\" \"(and a b)\") (\"b\" \"(or b (not a))\"))"))
    (define m2 (read-org-variable-mapping "((\"a\" . \"(and a b)\") (\"b\" . \"(or b (not a))\"))"))
    (define m3 (unorgv "((\"a\" . \"(and a b)\") (\"b\" . \"(or b (not a))\"))"))
    (check-equal? (hash-ref m1 'a) '(and a b))
    (check-equal? (hash-ref m2 'a) '(and a b))
    (check-equal? (hash-ref m3 'a) '(and a b))
    (check-equal? (hash-ref m1 'b) '(or b (not a)))
    (check-equal? (hash-ref m2 'b) '(or b (not a)))
    (check-equal? (hash-ref m3 'b) '(or b (not a)))))

;;; A synonym for read-org-variable-mapping.
(define unorgv read-org-variable-mapping)

;;; Typeset the graph via graphviz and display it.
(define dotit (compose display graphviz))

;;; Reads a list of symbols from a string.
(define (read-symbol-list str)
  (string->any (string-append "(" str ")")))

(module+ test
  (test-case "read-symbol-list"
    (check-equal? (read-symbol-list "a b c") '(a b c))))

;;; Removes the first and the last symbol of a given string.
;;;
;;; Useful for removing the parentheses in string representations of
;;; lists.
(define (drop-first-last str)
  (substring str 1 (- (string-length str) 1)))

(module+ test
  (test-case "drop-first-last"
    (check-equal? (drop-first-last "(a b)") "a b")))

;;; Converts a list of sets of symbols to a list of strings containing
;;; those symbols.
(define (list-sets->list-strings lst)
  (map (compose drop-first-last any->string set->list) lst))

(module+ test
  (test-case "list-sets->list-strings"
    (check-equal? (list-sets->list-strings (list (set 'x 'y) (set 'z) (set) (set 't)))
                  '("y x" "z" "" "t"))))

;;; Pretty-prints a set of sets of symbols.
;;;
;;; Typically used for pretty-printing the annotations on the edges of
;;; the state graph.
(define (pretty-print-set-sets ms)
  (string-join (for/list ([m ms]) (format "{~a}" (pretty-print-set m))) ""))

(module+ test
  (test-case "pretty-print-set-sets"
    (check-equal? (pretty-print-set-sets (set (set 'a 'b) (set 'c))) "{a b}{c}")))


;;; ==========================
;;; Additional graph utilities
;;; ==========================

;;; Apply a transformation to every vertex in the unweighted graph,
;;; return the new graph.  If the transformation function maps two
;;; vertices to the same values, these vertices will be merged in the
;;; resulting graph.  The transformation function may be called
;;; multiple times for the same vertex.
;;;
;;; This function does not rely on rename-vertex!, so it can be used
;;; to permute vertex labels.
(define (update-vertices/unweighted gr func)
  (unweighted-graph/directed
   (for/list ([e (in-edges gr)])
     (match-let ([(list u v) e])
       (list (func u) (func v))))))

(module+ test
  (test-case "update-vertices/unweighted"
    (define gr1 (directed-graph '((a b) (b c))))
    (define gr2 (undirected-graph '((a b) (b c))))
    (define dbl (λ (x) (let ([x-str (symbol->string x)])
                         (string->symbol (string-append x-str x-str)))))
    (define new-gr1 (update-vertices/unweighted gr1 dbl))
    (define new-gr2 (update-vertices/unweighted gr2 dbl))

    (check-false (has-vertex? new-gr1 'a))
    (check-true (has-vertex? new-gr1 'aa))
    (check-false (has-vertex? new-gr1 'b))
    (check-true (has-vertex? new-gr1 'bb))
    (check-false (has-vertex? new-gr1 'c))
    (check-true (has-vertex? new-gr1 'cc))
    (check-true (has-edge? new-gr1 'aa 'bb))
    (check-true (has-edge? new-gr1 'bb 'cc))

    (check-true (has-edge? new-gr2 'aa 'bb))
    (check-true (has-edge? new-gr2 'bb 'aa))
    (check-true (has-edge? new-gr2 'bb 'cc))
    (check-true (has-edge? new-gr2 'cc 'bb))))

;;; Given a graph, apply a transformation v-func to every vertex label
;;; and, if the graph is a weighted graph, the transformation e-func
;;; to every edge label.  Both transformations default to identity
;;; functions.  If gr is an weighted graph, the result is a weighted
;;; graph.  If gr is an unweighted graph, the result is an unweighted
;;; graph.
(define (update-graph gr
                      #:v-func [v-func identity]
                      #:e-func [e-func identity])
  (define edges
    (for/list ([e (in-edges gr)])
      (match-let ([(list u v) e])
        (cond
          [(unweighted-graph? gr) (list (v-func u) (v-func v))]
          [else (list (e-func (edge-weight gr u v))
                      (v-func u) (v-func v))]))))
  (cond
    [(unweighted-graph? gr) (unweighted-graph/directed edges)]
    [else
     (weighted-graph/directed edges)]))

(module+ test
  (test-case "update-graph"
    (define gr1 (directed-graph '((a b) (b c))))
    (define gr2 (undirected-graph '((a b) (b c))))
    (define dbl (λ (x) (let ([x-str (symbol->string x)])
                         (string->symbol (string-append x-str x-str)))))
    (define new-gr1-ug (update-graph gr1 #:v-func dbl))
    (define new-gr2-ug (update-graph gr2 #:v-func dbl))
    (define gr3 (weighted-graph/directed '((10 a b) (11 b c))))
    (define new-gr3 (update-graph gr3 #:v-func dbl #:e-func (λ (x) (* 2 x))))

    (check-false (has-vertex? new-gr1-ug 'a))
    (check-true (has-vertex? new-gr1-ug 'aa))
    (check-false (has-vertex? new-gr1-ug 'b))
    (check-true (has-vertex? new-gr1-ug 'bb))
    (check-false (has-vertex? new-gr1-ug 'c))
    (check-true (has-vertex? new-gr1-ug 'cc))
    (check-true (has-edge? new-gr1-ug 'aa 'bb))
    (check-true (has-edge? new-gr1-ug 'bb 'cc))

    (check-true (has-edge? new-gr2-ug 'aa 'bb))
    (check-true (has-edge? new-gr2-ug 'bb 'aa))
    (check-true (has-edge? new-gr2-ug 'bb 'cc))
    (check-true (has-edge? new-gr2-ug 'cc 'bb))

    (check-true (has-edge? new-gr3 'aa 'bb))
    (check-false (has-edge? new-gr3 'bb 'aa))
    (check-true (has-edge? new-gr3 'bb 'cc))
    (check-false (has-edge? new-gr3 'cc 'bb))
    (check-equal? (edge-weight new-gr3 'aa 'bb) 20)
    (check-equal? (edge-weight new-gr3 'bb 'cc) 22)))


;;; ===============
;;; Pretty printing
;;; ===============

;;; Pretty print a set by listing its elements in alphabetic order.
(define (pretty-print-set s)
  (string-join (sort (set-map s any->string) string<?)))

(module+ test
  (test-case "pretty-print-set"
    (check-equal? (pretty-print-set (set 'a 'b 1)) "1 a b")))


;;; ======================================
;;; Additional list and hash map utilities
;;; ======================================

;;; Collects labels for duplicate edges into a sets of labels.
;;;
;;; More precisely, given a list of edges and weights, produces a new
;;; list of edges without duplicates, and a list of lists of weights
;;; in which each element corresponds to the edge (the input is
;;; suitable for graph constructors).
(define (collect-by-key edges labels)
  (for/fold ([ht (make-immutable-hash)]
             #:result (values (hash-keys ht) (hash-values ht)))
            ([e edges] [l labels])
    (hash-update ht e (λ (ls) (cons l ls)) empty)))

(module+ test
  (test-case "collect-by-key"
    (define-values (e1 l1) (collect-by-key '((1 2) (1 3)) '(a b)))
    (define-values (e2 l2) (collect-by-key '((1 2) (1 2)) '(a b)))
    (check-equal? e1 '((1 2) (1 3))) (check-equal? l1 '((a) (b)))
    (check-equal? e2 '((1 2))) (check-equal? l2 '((b a)))))

;;; Like collect-by-key, but returns a list of sets of weights.
(define (collect-by-key/sets edges labels)
  (let-values ([(es ls) (collect-by-key edges labels)])
    (values es (map list->set ls))))

(module+ test
  (test-case "collect-by-key/sets"
    (define-values (e3 l3) (collect-by-key/sets '(a b a) '(1 2 1)))
    (check-equal? e3 '(a b)) (check-equal? l3 (list (set 1) (set 2)))))

;;; Converts the values of a hash table from lists to sets.
(define (ht-values/list->set ht)
  (for/hash ([(k v) (in-hash ht)])
    (values k (list->set v))))

(module+ test
  (test-case "ht-values/list->set"
    (check-equal? (ht-values/list->set #hash((a . (1 1))))
                  (hash 'a (set 1)))))

;;; Returns the key-value pairs of a given hash table in the order in
;;; which the hash table orders them for hash-map and hash-for-each.
(define (hash->list/ordered ht) (hash-map ht cons #t))

(module+ test
  (test-case "hash->list/ordered"
    (check-equal? (hash->list/ordered #hash((b . 1) (a . 1)))
                  '((a . 1) (b . 1)))))

;;; Given a list of lists, splits every single list at the given
;;; position, and then returns two lists: one consisting of the first
;;; halves, and the one consisting of the second halves.
(define (multi-split-at lsts pos)
  (for/fold ([lefts '()]
             [rights '()]
             #:result (values (reverse lefts) (reverse rights)))
            ([lst (in-list lsts)])
    (define-values (left right) (split-at lst pos))
    (values (cons left lefts) (cons right rights))))

(module+ test
  (test-case "multi-split-at"
    (define-values (l1 l2) (multi-split-at '((1 2 3) (a b c)) 2))
    (check-equal? l1 '((1 2) (a b))) (check-equal? l2 '((3) (c)))))

;;; Given a list of lists of the same length, transposes them.
;;;
;;; > (lists-transpose '((1 2) (a b)))
;;; '((1 a) (2 b))
;;;
;;; This function is essentially in-parallel, wrapped in a couple
;;; conversions.
(define lists-transpose
  (compose sequence->list
           in-values-sequence
           ((curry apply) in-parallel)))

(module+ test
  (test-case "lists-transpose"
    (check-equal? (lists-transpose '((1 2) (a b))) '((1 a) (2 b)))))

;;; Like hash-union, but computes the intersection of hash
;;; tables.
;;;
;;; The intersections are computed by functional update, intersecting
;;; the first hash table with every other one supplied, one by one.
;;; The new mappings are computed either using combine or combine/key.
;;;
;;; TODO: I submitted this function to the Racket repository.  Remove
;;; this function from this module when the update gets applied.
(define (hash-intersect
         #:combine [combine #f]
         #:combine/key [combine/key
                        (if combine
                            (λ (_ x y) (combine x y))
                            (error 'hash-intersect))]
         one . rest)
  (define common-keys (apply set-intersect (map hash-keys (cons one rest))))
  (define res (for/hash ([k (in-list common-keys)])
                (values k (hash-ref one k))))
  (for*/fold ([res res]) ([hm (in-list rest)]
                          [(k v) (in-hash res)])
    (hash-set res k (combine/key k v (hash-ref hm k)))))

(module+ test
  (test-case "hash-intersect"
    (define h1 #hash((a . 1) (b . 2) (c . 3)))
    (define h2 #hash((a . 4) (b . 5)))
    (define h3 #hash((a . 7) (c . 8)))
    (check-equal? (hash-intersect h1 h2 h3 #:combine +)
                  '#hash((a . 12)))
    (check-equal? (hash-intersect h1 h3 #:combine -)
                  '#hash((a . -6) (c . -5)))))

;;; =========
;;; Functions
;;; =========

;;; Returns #t if the function has fixed arity (i.e. if it does not
;;; take a variable number of arguments).
(define (procedure-fixed-arity? func)
  (match (procedure-arity func)
    [(arity-at-least _) #f] [arity #t]))

(module+ test
  (test-case "procedure-fixed-arity?"
    (check-true (procedure-fixed-arity? not))
    (check-false (procedure-fixed-arity? +))))


;;; ==========
;;; Randomness
;;; ==========

;;; Generates a stream of inexact random numbers.  The meaning of the
;;; arguments is the same as for the function random:
;;;
;;; (in-randoms k) — a sequence of random exact integers in the range
;;; 0 to k-1.
;;;
;;; (in-randoms min max) — a sequence of random exact integers the
;;; range min to max-1.
;;;
;;; (in-randoms) — a sequence of random inexact numbers between
;;; 0 and 1.
(define in-random
  (case-lambda
    [() (for/stream ([i (in-naturals)]) (random))]
    [(k) (for/stream ([i (in-naturals)]) (random k))]
    [(min max) (for/stream ([i (in-naturals)]) (random min max))]))

(module+ test
  (test-case "in-random"
    (random-seed 0)
    (check-equal? (stream->list (stream-take (in-random 100) 10))
                  '(85 65 20 40 89 45 54 38 26 62))
    (check-equal? (stream->list (stream-take (in-random 50 100) 10))
                  '(75 59 82 85 61 85 59 64 75 53))
    (check-equal? (stream->list (stream-take (in-random) 10))
                  '(0.1656109603231493
                    0.9680391127132195
                    0.051518813640790355
                    0.755901955353936
                    0.5923534604277275
                    0.5513340634474264
                    0.7022057040731392
                    0.48375400938578744
                    0.7538961707172924
                    0.01828428516237329))))


;;; ===========================
;;; Additional stream utilities
;;; ===========================


;;; Returns the Cartesian product of the given streams.  The result is
;;; a stream whose elements are the elements of the Cartesian product.
;;;
;;; The implementation is inspired from the implementation of
;;; cartesian-product in racket/list.
(define (cartesian-product/stream . ss)
  ;; Cartesian product of two streams, produces an improper pair.
  (define (cp-2 ss1 ss2)
    (for*/stream ([s1 (in-stream ss1)] [s2 (in-stream ss2)]) (cons s1 s2)))
  ;; Fold-right over the list of streams.  The value for the fold is a
  ;; 1-value stream containing the empty list, which makes all the
  ;; lists proper.
  (foldr cp-2 (sequence->stream (in-value (list))) ss))

(module+ test
  (test-case "cartesian-product/stream"
    (check-equal? (stream->list (cartesian-product/stream (in-range 3) (in-range 4 6) '(a b)))
                  '((0 4 a)
                    (0 4 b)
                    (0 5 a)
                    (0 5 b)
                    (1 4 a)
                    (1 4 b)
                    (1 5 a)
                    (1 5 b)
                    (2 4 a)
                    (2 4 b)
                    (2 5 a)
                    (2 5 b)))))


;;; ==================
;;; Boolean operations
;;; ==================

;;; Returns the n-th Cartesian power of the Boolean domain: {0,1}^n.
(define (boolean-power n) (apply cartesian-product (make-list n '(#f #t))))

(module+ test
  (test-case "boolean-power"
    (check-equal? (boolean-power 2) '((#f #f) (#f #t) (#t #f) (#t #t)))))

;;; Like boolean-power, but returns a stream whose elements the
;;; elements of the Cartesian power.
(define (boolean-power/stream n) (apply cartesian-product/stream (make-list n '(#f #t))))

(module+ test
  (test-case "boolean-power/stream"
    (check-equal? (stream->list (boolean-power/stream 2)) '((#f #f) (#f #t) (#t #f) (#t #t)))))

;;; Converts any non-#f value to 1 and #f to 0.
(define (any->01 x) (if x 1 0))

(module+ test
  (test-case "any->01"
    (check-equal? (any->01 #t) 1)
    (check-equal? (any->01 #f) 0)))

;;; Converts 0 to #f and 1 to #t
(define (01->boolean x)
  (case x [(0) #f] [else #t]))

(module+ test
  (test-case "01->boolean"
    (check-equal? (01->boolean 0) #f)
    (check-equal? (01->boolean 1) #t)))