import collections
try:
    from collections.abc import Mapping
except ImportError:
    from collections import Mapping # noqa

import itertools
import numpy as np

try: from future_builtins import zip
except ImportError: pass

import segyio.tools as tools

class Gather(object):
    """
    A gather is in this context the intersection of lines in a cube, i.e. all
    the offsets at some inline/crossline intersection. The primary data type is
    numpy.ndarray, with dimensions depending on the range of offsets specified.

    Implements a dict_like lookup with the line and offset numbers (labels),
    not 0-based indices.

    Notes
    -----
    .. versionadded:: 1.1
    """

    def __init__(self, trace, iline, xline, offsets):
        # cache constructed modes for performance
        self.trace     = trace
        self.iline     = iline
        self.xline     = xline
        self.offsets   = offsets

    def __getitem__(self, index):
        """gather[i, x, o], gather[:,:,:]

        Get the gather or range of gathers, defined as offsets intersection
        between an in- and a crossline. Also works on post-stack files (with
        only 1 offset), although it is less useful in those cases.

        If offsets are omitted, the default is all offsets.

        A group of offsets is always returned as an offset-by-samples
        numpy.ndarray. If either inline, crossline, or both, are slices, a
        generator of such ndarrays are returned.

        If the slice of offsets misses all offsets, a special, empty ndarray is
        returned.

        Parameters
        ----------
        i : int or slice
            inline
        x : int or slice
            crossline
        o : int or slice
            offsets (default is :)

        Returns
        -------
        gather : numpy.ndarray or generator of numpy.ndarray

        Notes
        -----
        .. versionadded:: 1.1

        Examples
        --------
        Read one offset at an intersection:

        >>> gather[200, 241, 25] # returns same shape as trace

        Read all offsets at an intersection:

        >>> gather[200, 241, :] # returns offsets x samples ndarray
        >>> # If no offset is specified, this is implicitly (:)
        >>> gather[200, 241, :] == gather[200, 241]

        All offsets for a set of ilines, intersecting one crossline:

        >>> gather[200:300, 241, :] == gather[200:300, 241]

        Some offsets for a set of ilines, interescting one crossline:

        >>> gather[200:300, 241, 10:25:5]

        Some offsets for a set of ilines and xlines. This effectively yields a
        subcube:

        >>> f.gather[200:300, 241:248, 1:10]
        """
        if len(index) < 3:
            index = (index[0], index[1], None)

        il, xl, off = index

        if off is None and len(self.offsets) == 1:
            off = self.offsets[0]

        # if offset isn't specified, default to all, [:]
        off = off or slice(None)

        def isslice(x): return isinstance(x, slice)

        # gather[int,int,int]
        if not any(map(isslice, [il, xl, off])):
            o = self.iline.offsets[off]
            i = o + self.iline.heads[il] + self.xline.heads[xl]
            return self.trace[i]

        offs = off if isslice(off) else slice(off, off+1, 1)

        xs = list(filter(self.offsets.__contains__,
                    range(*offs.indices(self.offsets[-1]+1))))

        empty = np.empty(0, dtype = self.trace.dtype)
        # gather[int,int,:]
        if not any(map(isslice, [il, xl])):
            if len(xs) == 0: return empty
            i = self.iline.heads[il] + self.xline.heads[xl]
            return tools.collect(self.trace[i + self.iline.offsets[x]] for x in xs)

        # gather[:,:,:], gather[int,:,:], gather[:,int,:]
        # gather[:,:,int] etc
        def gen():
            # precomputed xline number -> xline offset into the iline
            xlinds = { xlno: i for i, xlno in enumerate(self.xline.keys()) }

            # ranges over gathers are done on a by-line basis so lines can be
            # buffered, and traces can be read from the iline. This is the
            # least efficient when there are very few traces read per inline,
            # but huge savings with larger subcubes
            last_il = self.iline.keys()[-1] + 1
            last_xl = self.xline.keys()[-1] + 1

            il_slice = il if isslice(il) else slice(il, il+1)
            xl_slice = xl if isslice(xl) else slice(xl, xl+1)

            # the slice could still be defaulted (:), in which case this starts
            # at zero. the lookups are guarded by try-excepts so it won't fail,
            # but it's unnecessary to chck all keys up until the first xline
            # because that will never be a hit anyway
            if il_slice.start is None:
                start = self.iline.keys()[0]
                il_slice = slice(start, il_slice.stop, il_slice.step)

            if xl_slice.start is None:
                start = self.xline.keys()[0]
                xl_slice = slice(start, xl_slice.stop, xl_slice.step)

            il_range = range(*il_slice.indices(last_il))
            xl_range = range(*xl_slice.indices(last_xl))

            # the try-except-else is essentially a filter on in/xl keys, but
            # delegates the work (and decision) to the iline and xline modes
            if not isslice(off):
                for iline in self.iline[il_slice, off]:
                    for xlno in xl_range:
                        try:
                            xind = xlinds[xlno]
                        except KeyError:
                            pass
                        else:
                            yield iline[xind]

                return

            if len(xs) == 0:
                for _, _ in itertools.product(il_range, xl_range): yield empty
                return

            for ilno in il_range:
                iline = tools.collect(self.iline[ilno, off])
                for x in xl_range:
                    try:
                        xind = xlinds[x]
                    except KeyError:
                        pass
                    else:
                        yield iline[:, xind]

        return gen()

class Group(object):
    """
    The inner representation of the Groups abstraction provided by Group.

    A collection of trace indices that have identical `key`.

    Notes
    -----
    .. versionadded:: 1.9
    """
    def __init__(self, key, parent, index):
        self.parent = parent
        self.index = index
        self.key = key

    @property
    def header(self):
        """
        A generator of the the read-only headers in this group

        Returns
        -------
        headers : iterator of Header

        Notes
        -----
        The generator respects the order of the index - to iterate over headers
        in a different order, the index attribute can be re-organised.

        .. versionadded:: 1.9
        """
        source = self.parent.header
        for i in self.index:
            yield source[i]

    @property
    def trace(self):
        """
        A generator of the the read-only traces in this group

        Returns
        -------
        traces : iterator of Trace

        Notes
        -----
        The generator respects the order of the index - to iterate over headers
        in a different order, the index attribute can be re-organised.

        .. versionadded:: 1.9
        """
        source = self.parent.trace
        for i in self.index:
            yield source[i]

    def sort(self, fields):
        """
        Sort the traces in the group, obeying the `fields` order of
        most-to-least significant word.
        """
        # TODO: examples

        headers = [dict(self.parent.header[i]) for i in self.index]
        index = list(zip(headers, self.index))
        # sorting is stable, so sort the whole set by field, applied in the
        # reverse order:
        for field in reversed(fields):
            index.sort(key = lambda x: x[0][field])

        # strip off all the headers
        index = [i for _, i in index]
        self.index = index

class Groups(Mapping):
    """
    The Groups implements the dict interface, grouping all traces that match a
    given `fingerprint`. The fingerprint is a signature derived from a set of
    trace header words, called a `key`.

    Consider a file with five traces, and some selected header words::

        0: {offset: 1, fldr: 1}
        1: {offset: 1, fldr: 2}
        2: {offset: 1, fldr: 1}
        3: {offset: 2, fldr: 1}
        4: {offset: 1, fldr: 2}

    With key = (offset, fldr), there are 3 groups::

        {offset: 1, fldr: 1 }: [0, 2]
        {offset: 1, fldr: 2 }: [1, 4]
        {offset: 2, fldr: 1 }: [3]

    With a key = offset, there are 2 groups::

        {offset: 1}: [0, 1, 2, 4]
        {offset: 2}: [3]

    The Groups class is intended to easily process files without the rigid
    in/crossline structure of iline/xline/gather, but where there is sufficient
    structure in the headers. This is common for some types of pre-stack data,
    shot gather data etc.

    Notes
    -----
    .. versionadded:: 1.9
    """
    # TODO: only group in range of traces?
    # TODO: cache header dicts?
    def __init__(self, trace, header, key):
        bins = collections.OrderedDict()
        for i, h in enumerate(header[:]):
            k = self.fingerprint(h[key])
            if k in bins:
                bins[k].append(i)
            else:
                bins[k] = [i]

        self.trace = trace
        self.header = header
        self.key = key
        self.bins = bins

    @staticmethod
    def normalize_keys(items):
        """
        Normalize the key representation to integers, so that they're hashable,
        even when a key is built with enumerators.

        This function is intended for internal use, and provides the mapping
        from accepted key representation to a canonical key.

        Parameters
        ----------
        items : iterator of (int_like, array_like)

        Returns
        -------
        items : generator of (int, array_like)

        Warnings
        --------
        This function provides no guarantees for value and type compatibility,
        even between minor versions.

        Notes
        -----
        .. versionadded:: 1.9
        """
        return ((int(k), v) for k, v in items)

    @staticmethod
    def fingerprint(key):
        """
        Compute a hashable fingerprint for a key. This function is intended for
        internal use. Relies on normalize_keys for transforming keys to
        canonical form. The output of this function is used for the group ->
        index mapping.

        Parameters
        ----------
        key : int_like or dict of {int_like: int} or iterable of (int_like,int)

        Returns
        -------
        key
            A normalized canonical representation of key

        Warnings
        --------
        This function provides no guarantees for value and type compatibility,
        even between minor versions.

        Notes
        -----
        .. versionadded:: 1.9
        """
        try:
            return int(key)
        except TypeError:
            pass

        try:
            items = key.items()
        except AttributeError:
            items = iter(key)

        # map k -> tracefield -> int
        return frozenset(Groups.normalize_keys(items))

    def __len__(self):
        """x.__len__() <==> len(x)"""
        return len(self.bins)

    def __contains__(self, key):
        """x.__len__() <==> len(x)"""
        return self.fingerprint(key) in self.bins

    def __getitem__(self, key):
        """g[key]

        Read the group associated with key.

        Key can be any informal mapping between a header word (TraceField, su
        header words, or raw integers) and a value.

        Parameters
        ----------
        key

        Returns
        -------
        group : Group

        Notes
        -----
        .. versionadded:: 1.9

        Examples
        --------

        Group on FieldRecord, and get the group FieldRecord == 5:

        >>> fr = segyio.TraceField.FieldRecord
        >>> records = f.group(fr)
        >>> record5 = records[5]
        """
        key = self.fingerprint(key)
        return Group(key, self, self.bins[key])

    def values(self):
        for key, index in self.bins.items():
            yield Group(key, self, index)

    def items(self):
        for key, index in self.bins.items():
            yield key, Group(key, self, index)

    def __iter__(self):
        return self.bins.keys()

    def sort(self, fields):
        """
        Reorganise the indices in all groups by fields
        """
        bins = collections.OrderedDict()

        for key, index in self.bins.items():
            g = Group(key, self, index)
            g.sort(fields)
            bins[key] = g.index

        self.bins = bins
