namespace Bernsrite.Nessos.Streams

open Nessos.Streams

module Stream =

    /// <summary>Applies a key-generating function to each element of the input stream and yields a stream of unique keys and a stream of all elements that have each key.</summary>
    /// <param name="projection">A function to transform items of the input stream into comparable keys.</param>
    /// <param name="source">The input stream.</param>
    /// <returns>A stream of tuples where each tuple contains the unique key and a stream of all the elements that match the key.</returns>    
    let groupByStreams (projection : 'elem -> 'key) source =
        source
            |> Stream.groupBy projection
            |> Stream.map (fun (key, grouping) ->
                key, Stream.ofSeq grouping)

    /// <summary>Concatenates a stream of streams.</summary>
    /// <param name="streams">The stream of streams to concatenate.</param>
    /// <returns>The concatenated stream.</returns>
    let concatStreams streams =
        streams
            |> Stream.toSeq
            |> Stream.concat

    /// Wraps the two given streams as a single concatenated stream.
    let append stream1 stream2 =
        Stream.concat [stream1; stream2]

    /// Returns the average of the elements in the stream.
    let inline average stream =
        stream
            |> Stream.toSeq
            |> Seq.average

    /// Returns the average of the results generated by applying
    /// the function to each element of the stream.
    let inline averageBy projection stream =
        stream
            |> Stream.toSeq
            |> Seq.averageBy projection

    /// Applies the given function to each element of the stream
    /// and concatenates all the results.
    let collect mapping stream =
        stream
            |> Stream.map mapping
            |> concatStreams

    /// <summary>Tests if the stream contains the specified element.</summary>
    /// <param name="value">The value to locate in the input stream.</param>
    /// <param name="stream">The input stream.</param>
    let contains value stream =
        stream
            |> Stream.tryFind ((=) value)
            |> Option.isSome

    /// Returns a sequence that contains no duplicate entries according to
    /// generic hash and equality comparisons on the entries. If an element
    /// occurs multiple times in the sequence then the later occurrences are
    /// discarded.
    let distinct stream =
        stream
            |> Stream.toSeq
            |> Seq.distinct
            |> Stream.ofSeq

    /// Returns the only element of the stream.
    let exactlyOne stream =
        stream
            |> Stream.toSeq
            |> Seq.exactlyOne

    /// <summary>Generates a new stream which contains successive elements generated
    /// by calling the given function, up to the given count. The function is passed
    /// the index of the item being generated.</summary>
    /// <param name="count">The number of items to generate for the stream.</param>
    /// <param name="initializer">A function that generates an item in the sequence from a given index.</param>
    let init count initializer =
        Seq.init count initializer
            |> Stream.ofSeq

    /// <summary>Applies the given function to each element of the stream. The integer passed to the
    /// function indicates the index of element.</summary>
    /// <param name="action">A function to apply to each element of the stream that can also access the current index.</param>
    /// <param name="stream">The input stream.</param>
    let iteri action stream =
        stream
            |> Stream.mapi (fun i elem -> i, elem)
            |> Stream.iter (fun (i, elem) -> action i elem)

    /// Returns the greatest of all elements of the stream.
    let max stream =
        stream
            |> Stream.maxBy id

    /// Returns the lowest of all elements of the stream.
    let min stream =
        stream
            |> Stream.minBy id

    /// Returns a stream of each element in the input stream and its predecessor,
    /// with the exception of the first element which is only returned as the
    /// predecessor of the second element.
    let pairwise stream =
        stream
            |> Stream.toSeq
            |> Seq.pairwise
            |> Stream.ofSeq

    /// <summary>Returns a stream that, when iterated, skips elements of the underlying stream while the 
    /// given predicate returns True, and then yields the remaining elements of the stream.</summary>
    /// <param name="predicate">A function that evaluates an element of the stream to a boolean value.</param>
    /// <param name="stream">The input stream.</param>
    let skipWhile predicate stream =
        stream
            |> Stream.toSeq
            |> Seq.skipWhile predicate
            |> Stream.ofSeq

    /// Returns a stream ordered by its values.
    let sort stream =
        stream
            |> Stream.sortBy id

    /// <summary>Returns the sum of the results generated by applying the function to each element of the stream.</summary>
    /// <param name="projection">A function to transform items from the input stream into the type that will be summed.</param>
    /// <param name="stream">The input stream.</param>
    let sumBy projection stream =
        stream
            |> Stream.minBy projection
            |> Stream.sum

    /// Returns a stream that skips 1 element of the underlying stream and
    /// then contains the remaining elements of the stream.
    let tail stream =
        stream
            |> Stream.skip 1

    /// <summary>Returns a stream that contains at most the first N elements of the given stream.</summary>
    /// <param name="count">The maximum number of items to include.</param>
    /// <param name="stream">The input stream.</param>
    let truncate n stream =
        stream
            |> Stream.take n   // current implementation of `take` actually implements `truncate` functionality

    /// Filters the elements of the input stream.
    let where predicate stream =
        stream
            |> Stream.filter predicate

    /// <summary>Combines the two streams into a stream of pairs.</summary>
    /// <param name="stream1">The first input stream.</param>
    /// <param name="stream2">The second input stream.</param>
    let zip stream1 stream2 =
        Stream.zipWith
            (fun elem1 elem2 -> elem1, elem2)
            stream1
            stream2

/// Builder type for streams.
type StreamBuilder() =

    /// Combines two streams.
    member __.Combine(stream1, stream2) =
        Stream.concat [stream1; stream2]

    /// Executes the given function.
    member __.Delay(f) =
        f()

    /// Iterates the given items, generating a stream.
    member __.For(items, f ) =
        items
            |> Seq.map f
            |> Stream.concat

    /// Yields the given item as a stream.
    member __.Yield(item) =
        Stream.ofSeq [item]

    /// Yields the given items as a stream.
    member __.YieldFrom(items) =
        Stream.ofSeq items

    /// The empty stream.
    member __.Zero() =
        Stream.empty

[<AutoOpen>]
module Builder =

    /// Builder object for streams.
    let stream = new StreamBuilder()
