"""The Stream module contains the class implementing streams, or lazily
computed recursive lists, for CS61A at UC Berkeley.
"""


class Stream(object):
    """A lazily computed recursive list."""
    def __init__(self, first, compute_rest, empty=False):
        self.first = first
        self._compute_rest = compute_rest
        self.empty = empty
        self._rest = None
        self._computed = False
        
    @property
    def rest(self):
        """Return the rest of the stream, computing it if necessary."""
        assert not self.empty, 'Empty streams have no rest.'
        if not self._computed:
            self._rest = self._compute_rest()
            self._computed = True
        return self._rest
        
    def __repr__(self):
        if self.empty:
            return '<empty stream>'
        return 'Stream({0}, <compute_rest>)'.format(repr(self.first))


# The empty stream.
Stream.the_empty_stream = Stream(None, None, True)


# Utility functions for streams
def show_stream(strm, count=10):
    """Expand up to count (default 10) items into the stream and return the
    string representation of the stream.

    >>> ones = Stream(1, lambda: ones)
    >>> zeroes = Stream(0, lambda: zeroes)
    >>> show_stream(ones)
    '1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...'
    >>> show_stream(zeroes, count=5)
    '0, 0, 0, 0, 0, ...'
    """
    
    result = ''
    furthest_part = strm
    for _ in range(count):
        str_representation = '\'' + show_stream(furthest_part.first, count) + '\'' \
                             if isinstance(furthest_part.first, Stream) \
                             else str(furthest_part.first)
        result += '{0}'.format(str_representation)
        if furthest_part.rest is Stream.the_empty_stream:
            return result
        furthest_part = furthest_part.rest
        result += ', '
    return result + '...'


def add_streams(s1, s2):
    """Returns the stream where each element is the sum of the
    elements at the corresponding positions in the two streams.
    The result is as long as the shorter of the two streams.
    
    >>> naturals = integers_starting_from(1)
    >>> ones = make_ones()
    >>> sum_stream = add_streams(naturals, ones)
    >>> sum_stream.first
    2
    >>> sum_stream = add_streams(naturals, integers_starting_from(5))
    >>> sum_stream.rest.rest.first
    10
    """
    if s1.empty or s2.empty:
        return Stream.the_empty_stream
    return Stream(s1.first + s2.first, lambda: add_streams(s1.rest, s2.rest))
    
    
def mul_streams(s1, s2):
    """Returns the stream where each element is the product of the
    elements at the corresponding positions in the two streams.
    The result is as long as the shorter of the two streams.
    
    >>> naturals = integers_starting_from(1)
    >>> ones = make_ones()
    >>> prod_stream = mul_streams(naturals, ones)
    >>> prod_stream.rest.first
    2
    >>> show_stream(prod_stream)
    '1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...'
    >>> prod_stream = mul_streams(naturals, integers_starting_from(5))
    >>> prod_stream.rest.rest.first
    21
    >>> show_stream(prod_stream)
    '5, 12, 21, 32, 45, 60, 77, 96, 117, 140, ...'
    """
    if s1.empty or s2.empty:
        return Stream.the_empty_stream
    return Stream(s1.first * s2.first, lambda: mul_streams(s1.rest, s2.rest))
    
    
def stream_map(fn, *strms):
    """Map a function over a (collection of) stream(s).

    >>> from operator import add
    >>> ones = Stream(1, lambda: ones)
    >>> ints = Stream(1, lambda: stream_map(add, ints, ones))
    >>> show_stream(ints)
    '1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...'
    >>> show_stream(stream_map(lambda x: x * 2, ints))
    '2, 4, 6, 8, 10, 12, 14, 16, 18, 20, ...'
    """
    if len(strms) == 0 or len([s for s in strms if s.empty]) > 0:
        return Stream.the_empty_stream
    args = [s.first for s in strms]
    return Stream(fn(*args), lambda: stream_map(fn, *[s.rest for s in strms]))


def stream_filter(pred, strm):
    """Filter items from a stream using the predicate pred.

    >>> from operator import add
    >>> odd = lambda x: x % 2 == 1
    >>> ones = Stream(1, lambda: ones)
    >>> ints = Stream(1, lambda: stream_map(add, ints, ones))
    >>> show_stream(stream_filter(odd, ints))
    '1, 3, 5, 7, 9, 11, 13, 15, 17, 19, ...'
    """
    if strm.empty:
        return Stream.the_empty_stream
    elif pred(strm.first):
        return Stream(strm.first, lambda: stream_filter(pred, strm.rest))
    return stream_filter(pred, strm.rest)


def stream_interleave(s1, s2):
    """Interleave the items from the two streams provided.
    
    >>> ones = make_ones()
    >>> naturals = integers_starting_from(1)
    >>> interleaved = stream_interleave(ones, naturals)
    >>> show_stream(interleaved)
    '1, 1, 1, 2, 1, 3, 1, 4, 1, 5, ...'
    """
    if s1.empty:
        return s2
    return Stream(s1.first, lambda: stream_interleave(s2, s1.rest))
    
    
def convert_to_stream(iterable):
    """Converts an iterable object to the corresponding stream.
    
    >>> my_list = [1, 6, 1, 8, 0, 3]
    >>> my_stream = convert_to_stream(my_list)
    >>> show_stream(my_stream)
    '1, 6, 1, 8, 0, 3'
    
    >>> def naturals():
    ...     curr = 1
    ...     while True:
    ...         yield curr
    ...         curr += 1
    ...
    >>> my_naturals = naturals()
    >>> my_naturals_stream = convert_to_stream(my_naturals)
    >>> show_stream(my_naturals_stream)
    '1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...'
    """
    iterator = iter(iterable)
    try:
        return Stream(next(iterator), lambda: convert_to_stream(iterator))
    except StopIteration:
        return Stream.the_empty_stream
    
    
# Common streams
def make_stream_of(n):
    """Returns an infinite stream of the same number n.
    
    >>> fives = make_stream_of(5)
    >>> fives.first
    5
    >>> fives.rest.rest.rest.first
    5
    >>> show_stream(fives)
    '5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...'
    """
    return Stream(n, lambda: make_stream_of(n))
    
    
def make_ones():
    """Returns an infinite stream of ones."""
    return make_stream_of(1)
    
    
def integers_starting_from(n):
    """Returns a stream of integers starting from the number n.
    
    >>> naturals = integers_starting_from(1)
    >>> naturals.first
    1
    >>> naturals.rest.first
    2
    >>> naturals.rest.rest.rest.rest.first
    5
    >>> show_stream(naturals)
    '1, 2, 3, 4, 5, 6, 7, 8, 9, 10, ...'
    """
    return Stream(n, lambda: integers_starting_from(n + 1))
    
    
def rand_stream(start=1, end=100):
    """Returns a stream of random integers, where each integer is in
    the range from start to end.
    """
    from random import randint
    return Stream(randint(start, end), lambda: rand_stream(start, end))