(module stream mzscheme
  (require (lib "stream.ss" "srfi" "40")
           (lib "list.ss" "srfi" "1")
           (lib "etc.ss"))

  ;; ===========================================================================
  ;; STREAM CONSTRUCTORS
  ;; ===========================================================================

  ;; stream-xcons : a a -> (streamof a)
  (define (stream-xcons d a)
    (stream-cons a d))

  ;; stream-cons* : a ... (streamof a) -> (streamof a)
  (define (stream-cons* elt0 . elts)
    (if (null? elts)
        elt0
        (stream-cons elt0 (apply stream-cons* (car elts) (cdr elts)))))

  ;; make-stream : (union nat +inf.0) * [a] -> (streamof a)
  (define make-stream
    (opt-lambda (n [fill #f])
      (stream-delay
       (if (zero? n)
           stream-null
           (stream-cons fill (make-stream (sub1 n) fill))))))

  ;; stream-tabulate : nat * (nat -> a) -> (streamof a)
  (define (stream-tabulate n init-proc)
    (recur stream-tabulate ([i 0])
      (stream-delay
       (if (= i n)
           stream-null
           (stream-cons (init-proc i) (stream-tabulate (add1 i)))))))

  ;; stream-copy : (streamof a) -> (streamof a)
  (define (stream-copy stream)
    (stream-delay
     (if (stream-null? stream)
         stream-null
         (stream-cons (stream-car stream) (stream-copy (stream-cdr stream))))))

  ;; circular-stream : a ... -> (streamof a)
  (define (circular-stream . elts)
    (letrec ([result (recur circular-stream ([prefix elts])
                       (stream-delay
                        (if (null? prefix)
                            result
                            (stream-cons (car prefix) (circular-stream (cdr prefix))))))])
      result))

  ;; stream-iota : [(union nat +inf.0)] * [nat] * [nat] -> (streamof nat)
  (define stream-iota
    (opt-lambda ([count +inf.0] [start 0] [step 1])
      (recur stream-iota ([i 0] [num start])
        (if (= i count)
            stream-null
            (stream-cons num (stream-iota (add1 i) (+ num step)))))))

  ;; list->stream : (listof a) -> (streamof a)
  (define (list->stream list)
    (apply stream list))

  ;; thunk->stream : (-> a) * [(a -> boolean)] -> (streamof a)
  (define thunk->stream
    (opt-lambda (proc [done? not])
      (recur thunk->stream ()
        (stream-delay
         (let ([result (proc)])
           (if (done? result)
               stream-null
               (stream-cons result (thunk->stream))))))))

  ;; ===========================================================================
  ;; SELECTORS
  ;; ===========================================================================

  (define (stream-first stream)
    (stream-car stream))
  (define (stream-second stream)
    (stream-car (stream-cdr stream)))
  (define (stream-third stream)
    (stream-car (stream-cdr (stream-cdr stream))))
  (define (stream-fourth stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr stream)))))
  (define (stream-fifth stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr (stream-cdr stream))))))
  (define (stream-sixth stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr stream)))))))
  (define (stream-seventh stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr stream))))))))
  (define (stream-eighth stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr stream)))))))))
  (define (stream-ninth stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr stream))))))))))
  (define (stream-tenth stream)
    (stream-car (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr (stream-cdr stream)))))))))))

  (define (stream-take i stream)
    (stream-delay
     (if (or (zero? i) (stream-null? stream))
         stream-null
         (stream-cons (stream-car stream)
                      (stream-take (sub1 i) (stream-cdr stream))))))

  (define (stream-drop i stream)
    (stream-delay
     (cond
       [(zero? i) stream]
       [(stream-null? stream) stream-null]
       [else (stream-drop (sub1 i) (stream-cdr stream))])))

  (define (stream-split-at i stream)
    (stream-delay
     (cond
       [(zero? i)
        (values stream-null stream)]
       [(stream-null? stream)
        (values stream-null stream-null)]
       [else
        (let-values ([(prefix suffix) (stream-split-at (sub1 i) (stream-cdr stream))])
          (values (stream-cons (stream-car stream) prefix) suffix))])))

  ;; ===========================================================================
  ;; MISCELLANEOUS
  ;; ===========================================================================

  ;; stream-append : (streamof a) ... -> (streamof a)
  (define (stream-append . streams)
    (cond
      [(null? streams) stream-null]
      [(null? (cdr streams)) (car streams)]
      [else (let ([stream1 (car streams)]
                  [streams (cdr streams)])
              (stream-delay
               (if (stream-null? stream1)
                   (apply stream-append streams)
                   (stream-cons (stream-car stream1)
                                (apply stream-append (stream-cdr stream1) streams)))))]))

  ;; stream-zip : (streamof a) ... -> (streamof (listof a))
  (define (stream-zip . streams)
    (stream-delay
     (if (any stream-null? streams)
         stream-null
         (stream-cons (map stream-car streams)
                      (apply stream-zip (map stream-cdr streams))))))

  ;; stream-filter-map : (a -> (optional b)) * (streamof a) -> (streamof b)
  (define (stream-filter-map proc stream)
    (stream-delay
     (cond
       [(stream-null? stream) stream-null]
       [(proc (stream-car stream))
        => (lambda (x)
             (stream-cons x (stream-filter-map proc (stream-cdr stream))))]
       [else (stream-filter-map proc (stream-cdr stream))])))

  ;; stream-append-map : (a -> (listof b)) * (streamof a) -> (streamof b)
  (define (stream-append-map proc stream)
    (stream-delay
     (if (stream-null? stream)
         stream-null
         (stream-append (proc (stream-car stream))
                        (stream-append-map proc (stream-cdr stream))))))

  ;; ===========================================================================
  ;; FILTERING AND PARTITIONING
  ;; ===========================================================================

  (define (stream-partition pred stream)
    (stream-delay
     (if (stream-null? stream)
         (values stream-null stream-null)
         (let ([x (stream-car stream)])
           (let-values ([(in out) (stream-partition pred (stream-cdr stream))])
             (if (pred x)
                 (values (stream-cons x in) out)
                 (values in (stream-cons x out))))))))

  (define (stream-remove pred stream)
    (stream-delay
     (cond
       [(stream-null? stream) stream-null]
       [(pred (stream-car stream)) (stream-remove pred (stream-cdr stream))]
       [else (stream-cons (stream-car stream) (stream-remove pred (stream-cdr stream)))])))

  ;; ===========================================================================
  ;; SEARCHING
  ;; ===========================================================================

  (define stream-member
    (opt-lambda (x stream [= equal?])
      (cond
        [(stream-null? stream) #f]
        [(= (stream-car stream) x) stream]
        [else (stream-member x (stream-cdr stream) =)])))

  (define (stream-memq x stream)
    (stream-member x stream eq?))

  (define (stream-memv x stream)
    (stream-member x stream eqv?))

  (define (stream-find pred stream)
    (cond
      [(stream-null? stream) #f]
      [(pred (stream-car stream)) (stream-car stream)]
      [else (stream-find pred (stream-cdr stream))]))

  (define (stream-find-tail pred stream)
    (cond
      [(stream-null? stream) #f]
      [(pred (stream-car stream)) stream]
      [else (stream-find-tail pred (stream-cdr stream))]))

  (define (stream-take-while pred stream)
    (stream-delay
     (if (or (stream-null? stream)
             (not (pred (stream-car stream))))
         stream-null
         (stream-cons (stream-car stream) (stream-take-while pred (stream-cdr stream))))))

  (define (stream-drop-while pred stream)
    (stream-delay
     (if (or (stream-null? stream)
             (not (pred (stream-car stream))))
         stream
         (stream-drop-while pred (stream-cdr stream)))))

  (define (stream-span pred stream)
    (stream-delay
     (if (or (stream-null? stream)
             (not (pred (stream-car stream))))
         (values stream-null stream)
         (let-values ([(prefix suffix) (stream-span pred (stream-cdr stream))])
           (values (stream-cons (stream-car stream) prefix) suffix)))))

  (define (stream-break pred stream)
    (stream-span (lambda (x) (not pred)) stream))

  (define (stream-any pred . streams)
    (cond
      [(any stream-null? streams) #f]
      [(apply pred (map stream-car streams))]
      [else (apply stream-any pred (map stream-cdr streams))]))

  (define (stream-every pred . streams)
    (cond
      [(any stream-null? streams) #t]
      [(apply pred (map stream-car streams))
       => (lambda (result)
            (let ([rests (apply stream-cdr streams)])
              (if (any stream-null? rests)
                  result
                  (apply stream-every pred rests))))]
      [else #f]))

  (define (stream-index pred . streams)
    (recur loop ([i 0] [streams streams])
      (cond
        [(any stream-null? streams) #f]
        [(apply pred (map stream-car streams)) i]
        [else (loop (add1 i) (map stream-cdr streams))])))

  ;; ===========================================================================
  ;; DELETION
  ;; ===========================================================================

  (define stream-delete
    (opt-lambda (x stream [= equal?])
      (stream-remove (lambda (y) (= x y)) stream)))

  ;; ===========================================================================
  ;; CONVERSIONS
  ;; ===========================================================================

  ;; stream->list : (streamof a) -> (listof a)
  (define (stream->list stream)
    (if (stream-null? stream)
        null
        (cons (stream-car stream)
              (stream->list (stream-cdr stream)))))

  ;; stream-ref : (streamof a) * nat -> a
  (define (stream-ref stream n)
    (cond
      [(stream-null? stream)
       (error 'stream-ref "expected non-null stream, received null")]
      [(zero? n)
       (stream-car stream)]
      [else
       (stream-ref (stream-cdr stream) (sub1 n))]))

  (provide (all-defined)))