"""
Collection of utilities to manipulate structured arrays.

Most of these functions were initially implemented by John Hunter for matplotlib.
They have been rewritten and extended for convenience.


"""


import itertools
from itertools import chain as iterchain, repeat as iterrepeat, izip as iterizip
import numpy as np
from numpy import ndarray, recarray
import numpy.ma as ma
from numpy.ma import MaskedArray
from numpy.ma.mrecords import MaskedRecords

from numpy.lib._iotools import _is_string_like

_check_fill_value = np.ma.core._check_fill_value

__all__ = ['append_fields',
           'drop_fields',
           'find_duplicates',
           'get_fieldstructure',
           'join_by',
           'merge_arrays',
           'rec_append_fields', 'rec_drop_fields', 'rec_join',
           'recursive_fill_fields', 'rename_fields',
           'stack_arrays',
           ]


def recursive_fill_fields(input, output):
    """
    Fills fields from output with fields from input,
    with support for nested structures.

    Parameters
    ----------
    input : ndarray
        Input array.
    output : ndarray
        Output array.

    Notes
    -----
    * `output` should be at least the same size as `input`

    Examples
    --------
    >>> a = np.array([(1, 10.), (2, 20.)], dtype=[('A', int), ('B', float)])
    >>> b = np.zeros((3,), dtype=a.dtype)
    >>> recursive_fill_fields(a, b)
    np.array([(1, 10.), (2, 20.), (0, 0.)], dtype=[('A', int), ('B', float)])

    """
    newdtype = output.dtype
    for field in newdtype.names:
        try:
            current = input[field]
        except ValueError:
            continue
        if current.dtype.names:
            recursive_fill_fields(current, output[field])
        else:
            output[field][:len(current)] = current
    return output



def get_names(adtype):
    """
    Returns the field names of the input datatype as a tuple.

    Parameters
    ----------
    adtype : dtype
        Input datatype

    Examples
    --------
    >>> get_names(np.empty((1,), dtype=int)) is None
    True
    >>> get_names(np.empty((1,), dtype=[('A',int), ('B', float)]))
    ('A', 'B')
    >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])])
    >>> get_names(adtype)
    ('a', ('b', ('ba', 'bb')))
    """
    listnames = []
    names = adtype.names
    for name in names:
        current = adtype[name]
        if current.names:
            listnames.append((name, tuple(get_names(current))))
        else:
            listnames.append(name)
    return tuple(listnames) or None


def get_names_flat(adtype):
    """
    Returns the field names of the input datatype as a tuple. Nested structure
    are flattend beforehand.

    Parameters
    ----------
    adtype : dtype
        Input datatype

    Examples
    --------
    >>> get_names_flat(np.empty((1,), dtype=int)) is None
    True
    >>> get_names_flat(np.empty((1,), dtype=[('A',int), ('B', float)]))
    ('A', 'B')
    >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])])
    >>> get_names_flat(adtype)
    ('a', 'b', 'ba', 'bb')
    """
    listnames = []
    names = adtype.names
    for name in names:
        listnames.append(name)
        current = adtype[name]
        if current.names:
            listnames.extend(get_names_flat(current))
    return tuple(listnames) or None


def flatten_descr(ndtype):
    """
    Flatten a structured data-type description.

    Examples
    --------
    >>> ndtype = np.dtype([('a', '<i4'), ('b', [('ba', '<f8'), ('bb', '<i4')])])
    >>> flatten_descr(ndtype)
    (('a', dtype('int32')), ('ba', dtype('float64')), ('bb', dtype('int32')))

    """
    names = ndtype.names
    if names is None:
        return ndtype.descr
    else:
        descr = []
        for field in names:
            (typ, _) = ndtype.fields[field]
            if typ.names:
                descr.extend(flatten_descr(typ))
            else:
                descr.append((field, typ))
        return tuple(descr)


def zip_descr(seqarrays, flatten=False):
    """
    Combine the dtype description of a series of arrays.

    Parameters
    ----------
    seqarrays : sequence of arrays
        Sequence of arrays
    flatten : {boolean}, optional
        Whether to collapse nested descriptions.
    """
    newdtype = []
    if flatten:
        for a in seqarrays:
            newdtype.extend(flatten_descr(a.dtype))
    else:
        for a in seqarrays:
            current = a.dtype
            names = current.names or ()
            if len(names) > 1:
                newdtype.append(('', current.descr))
            else:
                newdtype.extend(current.descr)
    return np.dtype(newdtype).descr


def get_fieldstructure(adtype, lastname=None, parents=None,):
    """
    Returns a dictionary with fields as keys and a list of parent fields as values.

    This function is used to simplify access to fields nested in other fields.

    Parameters
    ----------
    adtype : np.dtype
        Input datatype
    lastname : optional
        Last processed field name (used internally during recursion).
    parents : dictionary
        Dictionary of parent fields (used interbally during recursion).

    Examples
    --------
    >>> ndtype =  np.dtype([('A', int), 
    ...                     ('B', [('BA', int),
    ...                            ('BB', [('BBA', int), ('BBB', int)])])])
    >>> get_fieldstructure(ndtype)
    {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'],
     'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']}
    
    """
    if parents is None:
        parents = {}
    names = adtype.names
    for name in names:
        current = adtype[name]
        if current.names:
            if lastname:
                parents[name] = [lastname,]
            else:
                parents[name] = []
            parents.update(get_fieldstructure(current, name, parents))
        else:
            lastparent = [_ for _ in (parents.get(lastname, []) or [])]
            if lastparent:
#                if (lastparent[-1] != lastname):
                    lastparent.append(lastname)
            elif lastname:
                lastparent = [lastname,]
            parents[name] = lastparent or []
    return parents or None


def _izip_fields_flat(iterable):
    """
    Returns an iterator of concatenated fields from a sequence of arrays,
    collapsing any nested structure.
    """
    for element in iterable:
        if isinstance(element, np.void):
            for f in _izip_fields_flat(tuple(element)):
                yield f
        else:
            yield element


def _izip_fields(iterable):
    """
    Returns an iterator of concatenated fields from a sequence of arrays.
    """
    for element in iterable:
        if hasattr(element, '__iter__') and not isinstance(element, basestring):
            for f in _izip_fields(element):
                yield f
        elif isinstance(element, np.void) and len(tuple(element)) == 1:
            for f in _izip_fields(element):
                yield f
        else:
            yield element


def izip_records(seqarrays, fill_value=None, flatten=True):
    """
    Returns an iterator of concatenated items from a sequence of arrays.

    Parameters
    ----------
    seqarray : sequence of arrays
        Sequence of arrays.
    fill_value : {None, integer}
        Value used to pad shorter iterables.
    flatten : {True, False}, 
        Whether to 
    """
    # OK, that's a complete ripoff from Python2.6 itertools.izip_longest
    def sentinel(counter = ([fill_value]*(len(seqarrays)-1)).pop):
        "Yields the fill_value or raises IndexError"
        yield counter()
    #
    fillers = iterrepeat(fill_value)
    iters = [iterchain(it, sentinel(), fillers) for it in seqarrays] 
    # Should we flatten the items, or just use a nested approach
    if flatten:
        zipfunc = _izip_fields_flat
    else:
        zipfunc = _izip_fields
    #
    try:
        for tup in iterizip(*iters):
            yield tuple(zipfunc(tup))
    except IndexError:
        pass


def _fix_output(output, usemask=True, asrecarray=False):
    """
    Private function: return a recarray, a ndarray, a MaskedArray
    or a MaskedRecords depending on the input parameters
    """
    if not isinstance(output, MaskedArray):
        usemask = False
    if usemask:
        if asrecarray:
            output = output.view(MaskedRecords)
    else:
        output = ma.filled(output)
        if asrecarray:
            output = output.view(recarray)
    return output


def _fix_defaults(output, defaults=None):
    """
    Update the fill_value and masked data of `output`
    from the default given in a dictionary defaults.
    """
    names = output.dtype.names
    (data, mask, fill_value) = (output.data, output.mask, output.fill_value)
    for (k, v) in (defaults or {}).iteritems():
        if k in names:
            fill_value[k] = v
            data[k][mask[k]] = v
    return output


def merge_arrays(seqarrays,
                 fill_value=-1, flatten=False, usemask=True, asrecarray=False):
    """
    Merge arrays field by field.

    Parameters
    ----------
    seqarrays : sequence of ndarrays
        Sequence of arrays
    fill_value : {float}, optional
        Filling value used to pad missing data on the shorter arrays.
    flatten : {False, True}, optional
        Whether to collapse nested fields.
    usemask : {False, True}, optional
        Whether to return a masked array or not.
    asrecarray : {False, True}, optional
        Whether to return a recarray (MaskedRecords) or not.

    Examples
    --------
    >>> merge_arrays((np.array([1, 2]), np.array([10., 20., 30.])))
    masked_array(data = [(1, 10.0) (2, 20.0) (--, 30.0)],
          mask = [(False, False) (False, False) (True, False)],
          fill_value=(999999, 1e+20)
          dtype=[('f0', '<i4'), ('f1', '<f8')])
    >>> merge_arrays((np.array([1, 2]), np.array([10., 20., 30.])),
    ...              usemask=False)
    array(data = [(1, 10.0) (2, 20.0) (-1, 30.0)],
          dtype=[('f0', '<i4'), ('f1', '<f8')])
    >>> merge_arrays((np.array([1, 2]).view([('a', int)]),
                      np.array([10., 20., 30.])),
                     usemask=False, asrecarray=True)
    rec.array(data = [(1, 10.0) (2, 20.0) (-1, 30.0)],
              dtype=[('a', int), ('f1', '<f8')])
    """
    if (len(seqarrays) == 1):
        seqarrays = seqarrays[0]
    if isinstance(seqarrays, ndarray):
        seqdtype = seqarrays.dtype
        if (not flatten) or \
           (zip_descr((seqarrays,), flatten=True) == seqdtype.descr):
            seqarrays = seqarrays.ravel()
            if not seqdtype.names:
                seqarrays = seqarrays.view([('', seqdtype)])
            if usemask:
                if asrecarray:
                    return seqarrays.view(MaskedRecords)
                return seqarrays.view(MaskedArray)
            elif asrecarray:
                return seqarrays.view(recarray)
            return seqarrays
        else:
            seqarrays = (seqarrays,)
    # Get the dtype
    newdtype = zip_descr(seqarrays, flatten=flatten)
    # Get the data and the fill_value from each array
    seqdata = [ma.getdata(a.ravel()) for a in seqarrays]
    seqmask = [ma.getmaskarray(a).ravel() for a in seqarrays]
    fill_value = [_check_fill_value(fill_value, a.dtype) for a in seqdata]
    # Make an iterator from each array, padding w/ fill_values
    maxlength = max(len(a) for a in seqarrays)
    for (i, (a, m, fval)) in enumerate(zip(seqdata, seqmask, fill_value)):
        # Flatten the fill_values if there's only one field
        if isinstance(fval, (ndarray, np.void)):
            fmsk = ma.ones((1,), m.dtype)[0]
            if len(fval.dtype) == 1:
                fval = fval.item()[0]
                fmsk = True
            else:
                # fval and fmsk should be np.void objects
                fval = np.array([fval,], dtype=a.dtype)[0]
#                fmsk = np.array([fmsk,], dtype=m.dtype)[0]
        else:
            fmsk = True
        nbmissing = (maxlength-len(a))
        seqdata[i] = iterchain(a, [fval]*nbmissing)
        seqmask[i] = iterchain(m, [fmsk]*nbmissing)
    #
    data = izip_records(seqdata, flatten=flatten)
    data = tuple(data)
    if usemask:
        mask = izip_records(seqmask, fill_value=True, flatten=flatten)
        mask = tuple(mask)
        output = ma.array(np.fromiter(data, dtype=newdtype))
        output._mask[:] = list(mask)
        if asrecarray:
            output = output.view(MaskedRecords)
    else:
        output = np.fromiter(data, dtype=newdtype)
        if asrecarray:
            output = output.view(recarray)
    return output



def drop_fields(base, drop_names, usemask=True, asrecarray=False):
    """
    Return a new array with fields in `drop_names` dropped.

    Nested fields are supported.

    Parameters
    ----------
    base : array
        Input array
    drop_names : string or sequence
        String or sequence of strings corresponding to the names of the fields
        to drop.
    usemask : {False, True}, optional
        Whether to return a masked array or not.
    asrecarray : string or sequence
        Whether to return a recarray or a mrecarray (`asrecarray=True`) or
        a plain ndarray or masked array with flexible dtype (`asrecarray=False`)

    Examples
    --------
    >>> a = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
                     dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
    >>> drop_fields(a, 'a')
    array([((2.0, 3),), ((5.0, 6),)], 
          dtype=[('b', [('ba', '<f8'), ('bb', '<i4')])])
    >>> drop_fields(a, 'ba')
    array([(1, (3,)), (4, (6,))], 
          dtype=[('a', '<i4'), ('b', [('bb', '<i4')])])
    >>> drop_fields(a, ['ba', 'bb'])
    array([(1,), (4,)], 
          dtype=[('a', '<i4')])
    """
    if _is_string_like(drop_names):
        drop_names = [drop_names,]
    else:
        drop_names = set(drop_names)
    #
    def _drop_descr(ndtype, drop_names):
        names = ndtype.names
        newdtype = []
        for name in names:
            current = ndtype[name]
            if name in drop_names:
                continue
            if current.names:
                descr = _drop_descr(current, drop_names)
                if descr:
                    newdtype.append((name, descr))
            else:
                newdtype.append((name, current))
        return newdtype
    #
    newdtype = _drop_descr(base.dtype, drop_names)
    if not newdtype:
        return None
    #
    output = np.empty(base.shape, dtype=newdtype)
    output = recursive_fill_fields(base, output)
    return _fix_output(output, usemask=usemask, asrecarray=asrecarray)


def rec_drop_fields(base, drop_names):
    """
    Returns a new numpy.recarray with fields in `drop_names` dropped.
    """
    return drop_fields(base, drop_names, usemask=False, asrecarray=True)



def rename_fields(base, namemapper):
    """
    Rename the fields from a flexible-datatype ndarray or recarray.

    Nested fields are supported.

    Parameters
    ----------
    base : ndarray
        Input array whose fields must be modified.
    namemapper : dictionary
        Dictionary mapping old field names to their new version.

    Examples
    --------
    >>> a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))],
                      dtype=[('a', int),
                             ('b', [('ba', float), ('bb', (float, 2))])])
    >>> rename_fields(a, {'a':'A', 'bb':'BB'})
    array([(1, (2.0, 3)), (4, (5.0, 6))], 
          dtype=[('A', '<i4'), ('b', [('ba', '<f8'), ('BB', '<i4')])])

    """
    def _recursive_rename_fields(ndtype, namemapper):
        newdtype = []
        for name in ndtype.names:
            newname = namemapper.get(name, name)
            current = ndtype[name]
            if current.names:
                newdtype.append((newname,
                                 _recursive_rename_fields(current, namemapper)))
            else:
                newdtype.append((newname, current))
        return newdtype
    newdtype = _recursive_rename_fields(base.dtype, namemapper)
    return base.view(newdtype)


def append_fields(base, names, data=None, dtypes=None, 
                  fill_value=-1, usemask=True, asrecarray=False):