"""
Read and write wiggle tracks.

.. todo:: Now that we also include the end positions in the index, it is
    possible to do random jumps inside a region with some educated guessing.
    Perfect hits would not be possible, since the length of the lines is
    variable.

.. moduleauthor:: Martijn Vermaat <martijn@vermaat.name>

.. Licensed under the MIT license, see the LICENSE file.
"""


import sys

from .parse import LineType, create_state, parse
from .index import ReadError, index, write_index


def walk(track=sys.stdin, force_index=False):
    """
    Walk over the track and yield (region, position, value) tuples.

    The values are always of type `int` or `float`.

    :arg track: Wiggle track.
    :type track: file
    :arg force_index: Force creating an index if it does not yet exist.
    :type force_index: bool

    :return: Tuples of (region, position, value) per defined position.
    :rtype: generator(str, int, _)

    Example::

        >>> for x in walk():
        ...     x
        ...
        ('chr18', 34344, 629.0)
        ('chr18', 34345, 649.0)
        ('chr18', 34446, 657.0)
        ('chrM',  308,   520.0)
        ('chrM',  309,   519.0)

    .. todo:: Optionally give a list of regions to walk, in that order.
    """
    # Todo: Do something with browser and track lines.
    # Todo: Better exceptions.
    # Todo: Prettify this code.
    # Todo: Detect if index does not agree with track.
    region = None

    idx, _ = index(track, force=force_index)

    if idx is None:
        regions = [None]
    else:
        # Todo: Sort in a way that is compatible with existing wiggle tracks.
        #     Inspiration could be sorted BAM files. GATK requires these to be
        #     sorted according to the order in the reference file.
        regions = sorted(r for r in idx if r != '_all')

    for expected_region in regions:

        if expected_region is not None:
            track.seek(idx[expected_region]['start'])

        state = create_state()

        for line in track:
            line_type, data = parse(line, state)
            if line_type == LineType.REGION:
                region = data
                if expected_region is not None and region != expected_region:
                    break
            elif line_type == LineType.DATA:
                # Optimization: A `while` loop is faster than `for` and
                # `range`.
                i = 0
                while i < data.span:
                    yield region, data.position + i, data.value
                    i += 1

        # Todo: If there is no index yet, but we read the whole file, write
        #     the index we just built. But if we go this route, I think we
        #     should somehow merge wiggle.walk with index.index since they
        #     will be doing virtually the same thing.
        #
        #     if expected_region is None:
        #        write_index(idx, track)


def zip_(*walkers):
    """
    Walk over all tracks simultaneously and for each position yield the
    region, position and a list of values for each track, or `None` in case
    the track has no value on the position.

    .. note:: This assumes the order of regions is compatible over all
        walkers. If you are unsure if this is the case for your input wiggle
        tracks, use the :func:`walk` function with the `force_index` keyword
        argument.

    :arg walkers: List of generators yielding tuples of (region, position,
        value) per defined position.
    :type walkers: list(generator(str, int, _))

    :return: Tuples of (region, position, values) per defined position.
    :rtype: generator(str, int, list(_))

    Example::

        >>> for x in zip_(walk(open('a.wig')), walk(open('b.wig'))):
        ...     x
        ...
        ('18', 7, [29.0, None])
        ('18', 8, [49.0, None])
        ('18', 9, [None, 87.0])
        ('MT', 1, [20.0, None])
        ('MT', 2, [36.0, 92.0])
    """
    # We work with a list of lookahead items. If a walker has no more items,
    # we use None in the lookahead list.
    items = []
    for walker in walkers:
        try:
            items.append(next(walker))
        except StopIteration:
            items.append(None)

    # Regions seen so far.
    regions = set()
    previous_region = None

    while True:
        # If all lookahead items are None, we are done.
        if not any(items):
            break

        # Get the next position to yield.
        region, position = min(item[0:2] for item in items if item is not None)

        # Check region order compatibility.
        if region != previous_region:
            if region in regions:
                raise Exception('The order of regions is not compatible')
            regions.add(region)
            previous_region = region

        # Yield all values at this position.
        values = [item[2] if item is not None and item[0:2] == (region, position)
                  else None for item in items]
        yield region, position, values

        # Advance the lookahead list where we just yielded a value.
        for i, item in enumerate(items):
            if item is not None and item[0:2] == (region, position):
                try:
                    items[i] = next(walkers[i])
                except StopIteration:
                    items[i] = None


def fill(walker, regions=None, filler=None, only_edges=False):
    """
    Fill in undefined positions with `filler` (or `None`).

    :arg walker: Tuple of (region, position, value) per defined position.
    :type walker: generator(str, int, _)
    :arg regions: Dictionary with regions as keys and (start, stop) tuples as
        values. If not `None`, fill positions from start to stop (both
        including) in these regions. If `None`, fill positions in all
        regions between their first and last defined positions.
    :type regions: dict(str, (int, int))
    :arg filler: Value to use for filling undefined positions.
    :type filler: _
    :arg only_edges: Only fill the first and last of continuously undefined
        positions.
    :type only_edges: bool

    :return: Tuples of (region, position, value) per position where value is
        `filler` if it was not defined in the original walker.
    :rtype: generator(str, int, _)

    Example::

        >>> for x in walk(open('a.wig')):
        ...     x
        ...
        ('MT', 3, 29.0)
        ('MT', 5, 49.0)