#lang typed/racket

(require typed/graph typed-compose typed/racket/stream syntax/parse/define
         (for-syntax syntax/parse racket/syntax racket/list))

(provide
 Variable VariableMapping GeneralPair NonemptyListof

 assert-type for/first/typed for*/first/typed define/abstract/error
 relax-arg-type/any
 eval-with eval1-with auto-hash-ref/explicit auto-hash-ref/:
 lambda/: λ/: define/:
 extract-symbols any->string stringify-variable-mapping string->any
 handle-org-booleans map-sexp read-org-sexp unorg unstringify-pairs
 read-org-variable-mapping unorgv read-symbol-list drop-first-last
 list-sets->list-strings pretty-print-set pretty-print-set-sets
 update-vertices/unweighted update-graph dotit collect-by-key
 collect-by-key/sets ht-values/list->set hash->list/ordered hash-replace-keys/ordered
 multi-split-at lists-transpose append-lists in-random cartesian-product-2/stream
 cartesian-product/stream boolean-power boolean-power/stream any->01
 01->boolean)

(module+ test
  (require typed/rackunit))

(define-type Variable Symbol)
(define-type (VariableMapping A) (Immutable-HashTable Variable A))

(define-type (NonemptyListof a) (Pairof a (Listof a)))

(define-syntax-parse-rule (assert-type e:expr type:expr)
  (assert e (make-predicate type)))

(define-for-syntax (make-for/first/typed-variant folder)
  (syntax-parser
    #:literals (:)
    [(_ :
        ty:expr ; These should probably be more specific.
        clauses:expr
        c ...)
     #`(#,folder : ty
        ([result : ty #f])
        clauses
        #:break (not (equal? result #f))
        c ...)]))

(define-syntax for/first/typed (make-for/first/typed-variant 'for/fold))
(define-syntax for*/first/typed (make-for/first/typed-variant 'for*/fold))

(module+ test
  (test-case "for/first/typed, for/first/typed*"
    (check-equal? (for/first/typed : (Option Integer)
                                   ([i (in-range 1 10)]
                                    #:when (zero? (modulo i 5)))
                    (* i 3))
                  15)
    (check-equal? (for*/first/typed : (Option (Pairof Integer Integer))
                                    ([i (in-range 1 10)]
                                     [j (in-range 1 10)]
                                     #:when (> (+ i j) 5)
                                     #:when (even? i)
                                     #:when (even? j))
                    (cons i j))
                  '(2 . 4))))

(define-syntax-parser define/abstract/error
  [(_ (name:id args:id ...))
   #`(define/public (name args ...) (error 'name "abstract method"))])

(define-syntax-parse-rule (relax-arg-type/any name:id arg-type:expr)
  (λ ([x : Any]) (name (assert-type x arg-type))))

(: eval-with (-> (VariableMapping Any) Any AnyValues))
(define (eval-with ht expr)
  (parameterize ([current-namespace (make-base-namespace)])
    (for ([(x val) (in-hash ht)]) (namespace-set-variable-value! x val))
    (eval expr)))

(: eval1-with (-> (VariableMapping Any) Any Any))
(define (eval1-with ht expr)
  (call-with-values (λ () (eval-with ht expr))
                    (λ args (car args))))

(module+ test
  (test-case "eval-with"
    (check-equal? (eval1-with (hash 'a 1 'b 2) '(+ a b 1))
                  4)
    (define ht : (VariableMapping Integer) (hash 'a 1 'b 2))
    (define expr : Any '(+ a b 1))
    (check-equal? (eval1-with ht expr)
                  4)))

(define-syntax-parser auto-hash-ref/explicit
  [(_ (ht:id keys:id ...) body:expr)
   #`(let #,(for/list ([key (syntax->list #'(keys ...))])
              `[,key (hash-ref ,#'ht ',key)])
         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)))

;;; Helper functions for auto-hash-ref/:.
(begin-for-syntax
  (define (colon? s)
    (and (symbol? s)
         (eq? (string-ref (symbol->string s) 0) #\:)))
  (define (collect-colons datum)
    (cond [(colon? datum) (list datum)]
          [(list? datum)
           (remove-duplicates (flatten (for/list ([el (in-list datum)])
                                         (collect-colons el))))]
          [else '()]))
  (define (strip-colon s)
    (string->symbol (substring (symbol->string s) 1))))

(define-syntax (auto-hash-ref/: stx)
  (syntax-parse stx
    [(_ ht:id body:expr)
     (define colons (collect-colons (syntax->datum stx)))
     (define bindings (for/list ([key colons])
                        `[,key (hash-ref ,#'ht ',(strip-colon key))]))
     (with-syntax ([bindings-stx (datum->syntax #'ht bindings)])
       #'(let bindings-stx 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)))

(define-syntax-parser lambda/:
  [(_ body:expr)
   #:with ht (format-id #'body "ht")
   #'(lambda (ht) (auto-hash-ref/: ht body))]
  [(_ type:expr body:expr)
   #:with ht (format-id #'body "ht")
   #'(lambda ([ht : type]) (auto-hash-ref/: ht body))])

(define-syntax-parser λ/:
  [(_ body:expr) #'(lambda/: body)]
  [(_ type:expr body:expr) #'(lambda/: type body)])

(define-syntax-parser define/:
  [(_ name:id body:expr) #'(define name (λ/: body))]
  [(_ name:id type:expr body:expr) #'(define name (λ/: type body))])

(module+ test
  (test-case "lambda/:, λ/:, define/:"
    (define st : (HashTable Symbol Integer) (hash 'a 1 'b 2))
    (check-equal? ((lambda/: (+ :a :b)) st) 3)
    (check-equal? ((lambda/: (HashTable Symbol Integer) (+ :a :b)) st) 3)
    (check-equal? ((λ/: (HashTable Symbol Integer) (+ :a :b)) st) 3)

    (: f1 (-> (HashTable Symbol Integer) Integer))
    (define/: f1 (+ :a :b))
    (check-equal? (f1 st) 3)

    (define/: f2 (HashTable Symbol Integer) (+ :a :b))
    (check-equal? (f2 st) 3)))

(: extract-symbols (-> Any (Listof Symbol)))
(define (extract-symbols form)
  (: extract-rec (-> Any (Listof Any)))
  (define (extract-rec form)
    (match form
      [(? symbol?) (list form)]
      [(? list?)
       (flatten (for/list : (Listof Any)
                          ([x form])
                  (extract-symbols x)))]
      [else '()]))
  (assert-type (extract-rec form) (Listof Symbol)))

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

(: any->string (-> Any 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")))

(: stringify-variable-mapping (-> (VariableMapping Any) (VariableMapping String)))
(define (stringify-variable-mapping ht)
  (for/hash : (VariableMapping String)
            ([(key val) (in-hash 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)")))

(: string->any (-> String Any))
(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 sexps 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.
(: handle-org-booleans (-> Any Any))
(define/match (handle-org-booleans datum)
  [("#t") #t]
  [("#f") #f]
  [((? list?)) (map handle-org-booleans datum)]
  [ (_) datum])

(module+ test
  (test-case "handle-org-booleans"
    (check-equal? (handle-org-booleans "#t") #t)
    (check-equal? (handle-org-booleans "#f") #f)
    (check-equal? (handle-org-booleans '("#t" "#f")) '(#t #f))
    (check-equal? (handle-org-booleans "t") "t")))

(: map-sexp (-> (-> Any Any) Any Any))
(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 (λ (x)
                              (assert x number?)
                              (add1 x))
                            '(1 2 (4 10) 3))
                  '(2 3 (5 11) 4))))

(: read-org-sexp (-> String Any))
(define read-org-sexp
  (compose ((curry map-sexp) (match-lambda
                               [(and (? string?) str) (string->any str)]
                               [x x]))
           string->any))
(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? (unorg "(#t \"#t\" \"#t \" '(1 2 \"#f\"))")
                  '(#t #t #t '(1 2 #f)))))

(define-type (GeneralPair A B) (U (Pair A B) (List A B)))

(: unstringify-pairs (-> (Listof (GeneralPair String Any))
                         (Listof (GeneralPair Symbol Any))))
(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)
       (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)))))))

(: read-org-variable-mapping (-> String (VariableMapping Any)))
(define read-org-variable-mapping
  (multi-compose
   (λ ([pairs : (Listof (Pair Symbol Any))])
     (make-immutable-hash pairs))
   (λ (sexp)
     (assert-type sexp (Listof (GeneralPair String Any)))
     (unstringify-pairs sexp))
   string->any))

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

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

(: read-symbol-list (-> String (Listof Symbol)))
(define (read-symbol-list str)
  (assert-type (string->any (string-append "(" str ")")) (Listof Symbol)))

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

(: drop-first-last (-> String String))
(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")))

(: list-sets->list-strings (-> (Listof (Setof Any)) (Listof String)))
(define (list-sets->list-strings lst)
  (map (multi-compose drop-first-last
                      any->string
                      (λ ([x : (Setof Any)])
                        (set->list x))) lst))

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

(: pretty-print-set (-> (U (Setof Any) (Listof Any)) String))
(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")
    (check-equal? (pretty-print-set (list 'a 'b 1)) "1 a b")))

(: pretty-print-set-sets (-> (U (Setof (Setof Any))
                                (Listof (Setof Any)))
                             String))
(define (pretty-print-set-sets ms)
  (string-join (for/list ([m ms]) : (Listof String)
                         (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}")
    (check-equal? (pretty-print-set-sets (list (set 'a 'b) (set 'c))) "{a b}{c}")))

(define dotit (compose display graphviz))

(: update-vertices/unweighted (-> Graph (-> Any Any) Graph))
(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 : Any])
      (assert x symbol?)
      (define 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))))

(: update-graph (->* (Graph) (#:v-func (-> Any Any) #:e-func (-> Any Any)) Graph))
(define (update-graph gr #:v-func [v-func identity] #:e-func [e-func identity])
  (cond
    [(unweighted-graph? gr)
     (unweighted-graph/directed
      (for/list ([e (in-edges gr)]) : (Listof (List Any Any))
                (match-let ([(list u v) e])
                  (list (v-func u) (v-func v)))))]
    [else
     (weighted-graph/directed
      (for/list ([e (in-edges gr)]) : (Listof (List Any Any Any))
                (match-let ([(list u v) e])
                  (list (e-func (edge-weight gr u v))
                        (v-func u) (v-func v)))))]))

(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 : Any])
      (assert x symbol?)
      (define 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)
                                             (assert x number?)
                                             (* 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)))

(: collect-by-key (All (a b) (-> (Listof a) (Listof b)
                                 (Values (Listof a) (Listof (Listof b))))))
(define (collect-by-key keys vals)
  (for/fold ([ht : (HashTable a (Listof b))
                 (make-immutable-hash)]
             #:result (values (hash-keys ht) (hash-values ht)))
            ([e keys]
             [l vals])
    ((inst hash-update a (Listof b)) 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)))))

(: collect-by-key/sets (All (a b) (-> (Listof a) (Listof b)
                                      (Values (Listof a) (Listof (Setof b))))))
(define (collect-by-key/sets edges labels)
  (define-values (es ls) (collect-by-key edges labels))
  (values es ((inst map (Setof b) (Listof b)) 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 '(b a)) (check-equal? l3 (list (set 2) (set 1)))))

;;; Converts the values of a hash table from lists to sets.
(: ht-values/list->set (All (a b) (-> (HashTable a (Listof b)) (HashTable a (Setof b)))))
(define (ht-values/list->set ht)
  (for/hash ([(k v) (in-hash ht)]) : (HashTable a (Setof b))
            (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)))))

;; TODO: Remove after Typed Racket has caught up with Racket 8.4.
(: hash->list/ordered (All (a b) (-> (HashTable a b) (Listof (Pairof a b)))))
(define (hash->list/ordered ht)
  ((inst hash-map a b (Pairof a b)) ht cons #t))

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

(: hash-replace-keys/ordered (All (K1 K2 V) (-> (Immutable-HashTable K1 V) (Listof K2)
                                                (Immutable-HashTable K2 V))))
(define (hash-replace-keys/ordered ht new-keys)
  (make-immutable-hash (map (λ ([new-k : K2] [pair : (Pairof K1 V)])
                              (cons new-k (cdr pair)))
                            new-keys
                            (hash->list/ordered ht))))

(module+ test
  (test-case "hash-replace-keys/ordered"
    (check-equal? (hash-replace-keys/ordered (hash 'a 1 'b 2) '(x y))
                  '#hash((x . 1) (y . 2)))))

(: multi-split-at (All (a) (-> (Listof (Listof a)) Integer
                               (Values (Listof (Listof a)) (Listof (Listof a))))))
(define (multi-split-at lists pos)
  (for/fold ([lefts : (Listof (Listof a)) '()]
             [rights : (Listof (Listof a)) '()]
             #:result (values (reverse lefts) (reverse rights)))
            ([lst (in-list lists)])
    (define-values (left right) ((inst split-at a) 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)))))

;; https://racket.discourse.group/t/get-to-type-apply-in-parallel-lst/683
;;
;; Same thread: (apply ((curry map) list) lsts), however I don't
;; feel like typing this right now (2022-02-18).
(: lists-transpose (All (a ...) (-> (List (Listof a) ... a) (Listof (List a ... a)))))
(define (lists-transpose lists)
  (sequence->list (in-values-sequence (apply in-parallel lists))))

(module untyped racket
  (provide (contract-out [lists-transpose (-> (listof (listof any/c)) (listof (listof any/c)))]))
  (define (lists-transpose lists)
    (sequence->list (in-values-sequence (apply in-parallel lists)))))

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

(: append-lists (All (a) (-> (Listof (List (Listof a) (Listof a))) (Listof (Listof a)))))
(define (append-lists lsts)
  (for/list ([pr lsts])
    (append (car pr) (cadr pr))))

(module+ test
  (test-case "append-lists"
    (check-equal? (append-lists '(((1 2) (a b))
                                  ((3 4) (c d))))
                  '((1 2 a b)
                    (3 4 c d)))))

(: in-random (case->
              (-> (Sequenceof Flonum))
              (-> Integer (Sequenceof Nonnegative-Fixnum))
              (-> Integer Integer (Sequenceof Nonnegative-Fixnum))))
(define in-random
  (case-lambda
    [() (stream-cons
         (assert-type (random) Flonum)
         (in-random))]
    [(k) (stream-cons
          (assert-type (random k) Nonnegative-Fixnum)
          (in-random k))]
    [(min max) (stream-cons
                (assert-type (random min max) Nonnegative-Fixnum)
                (in-random min max))]))

(module+ test
  (test-case "in-random"
    (random-seed 1)
    (check-equal? (stream->list (stream-take (in-random 100) 10))
                  '(50 84 10 99 94 88 43 41 63 50))
    (check-equal? (stream->list (stream-take (in-random 50 100) 10))
                  '(57 98 82 83 61 53 73 82 50 80))
    (check-equal? (stream->list (stream-take (in-random) 10))
                  '(0.2718099186980313
                    0.7319496826374751
                    0.17365244033739616
                    0.5593031443038616
                    0.3345256691289459
                    0.9845704615094365
                    0.05753824253751768
                    0.22552976312818723
                    0.21646500425988832
                    0.15188352823997242))))

(: cartesian-product-2/stream (All (a b) (-> (Sequenceof a) (Sequenceof b) (Sequenceof (Pair a b)))))
(define (cartesian-product-2/stream s1 s2)
  (: cp2-store (All (a b) (-> (Sequenceof a) (Sequenceof b) (Sequenceof b)
                              (Sequenceof (Pair a b)))))
  ;; The recursive implementation using s2-store as an accumulator.
  ;; Main idea: combine the elements of s1 with the element of s2
  ;; until they are exhausted, then restart with the next element of
  ;; s1 and the original content of s2.
  (define (cp2-store s1 s2 s2-store)
    (cond
      [(stream-empty? s1) (stream)]
      [(stream-empty? s2) (cp2-store (stream-rest s1) s2-store s2-store)]
      [else
       (stream-cons (cons (stream-first s1) (stream-first s2))
                    (cp2-store s1 (stream-rest s2) s2-store))]))
  (cp2-store s1 s2 s2))

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

(: cartesian-product/stream (All (a) (-> (Listof (Sequenceof a)) (Sequenceof (Listof a)))))
(define (cartesian-product/stream ss)
  (for/foldr ([prod (stream (list))])
    ([s (in-list ss)])
    (cartesian-product-2/stream s prod)))

(module+ test
  (test-case "cartesian-product/stream"
    (check-equal? (stream->list (cartesian-product/stream '())) '(()))
    (check-equal? (stream->list (cartesian-product/stream '((a b c))))
                  '((a) (b) (c)))
    (check-equal? (stream->list (cartesian-product/stream (list (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-power (-> Integer (Listof (Listof Boolean))))
(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)))))

(: boolean-power/stream (-> Integer (Sequenceof (Listof Boolean))))
(define (boolean-power/stream n) (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)))))

(: any->01 (-> Any (U Zero One)))
(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)))

(: 01->boolean (-> (U Zero One) Boolean))
(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)))