# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See COPYING file distributed along with the PyMVPA package for the
#   copyright and license terms.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
"""NeuroImaging Markup Language (NIML) support.

Supports storing most typical values (samples, feature attributes, sample
attributes, dataset attributes) that are in a dataset in NIML format, as
long as these values are array-like.

Notes
-----
No support for 'sophisticated' values such as Mappers

.. versionadded:: 2.3.0

"""

__docformat__ = 'restructuredtext'

import numpy as np

import os

from mvpa2.support.nibabel import afni_niml_dset as niml_dset

from mvpa2.base.collections import SampleAttributesCollection, \
        FeatureAttributesCollection, DatasetAttributesCollection, \
        ArrayCollectable


from mvpa2.base import warning, debug, externals
from mvpa2.datasets.base import Dataset
from mvpa2.base import dataset

if externals.exists('h5py'):
    from mvpa2.base.hdf5 import h5save, h5load

_PYMVPA_PREFIX = 'PYMVPA'
_PYMVPA_SEP = '_'

def from_niml(dset, fa_labels=[], sa_labels=[], a_labels=[]):
    '''Convert a NIML dataset to a Dataset

    Parameters
    ----------
    dset: dict
        Dictionary with NIML key-value pairs, such as obtained from
        mvpa2.support.nibabel.afni_niml_dset.read()
    fa_labels: list
        Keys in dset that are enforced to be feature attributes
    sa_labels: list
        Keys in dset that are enforced to be sample attributes
    a_labels: list
        Keys in dset that are enforced to be dataset attributes

    Returns
    -------
    dataset: mvpa2.base.Dataset
        a PyMVPA Dataset
    '''

    # check for singleton element
    if type(dset) is list and len(dset) == 1:
        # recursive call
        return from_niml(dset[0])

    if not type(dset) is dict:
        raise ValueError("Expected a dict")

    if not 'data' in dset:
        raise ValueError("dset with no data?")

    data = dset['data']
    if len(data.shape) == 1:
        nfeatures = data.shape[0]
        nsamples = 1
    else:
        nfeatures, nsamples = data.shape

    # some labels have predefined destinations
    sa_labels_ = ['labels', 'stats', 'chunks', 'targets'] + sa_labels
    fa_labels_ = ['node_indices', 'center_ids'] + fa_labels
    a_labels_ = ['history'] + a_labels
    ignore_labels = ('data', 'dset_type')

    sa = SampleAttributesCollection(length=nsamples)
    fa = FeatureAttributesCollection(length=nfeatures)
    a = DatasetAttributesCollection()

    labels_collections = [(sa_labels_, sa),
                          (fa_labels_, fa),
                          (a_labels_, a)]

    infix2collection = {'sa':sa,
                      'fa':fa,
                      'a':a}

    infix2length = {'sa':nsamples, 'fa':nfeatures}

    for k, v in dset.iteritems():
        if k in ignore_labels:
            continue

        if k.startswith(_PYMVPA_PREFIX + _PYMVPA_SEP):
            # special PYVMPA field - do the proper conversion
            k_split = k.split(_PYMVPA_SEP)
            if len(k_split) > 2:
                infix = k_split[1].lower()
                collection = infix2collection.get(infix, None)
                if not collection is None:
                    short_k = _PYMVPA_SEP.join(k_split[2:])
                    expected_length = infix2length.get(infix, None)
                    if expected_length:
                        if isinstance(v, np.ndarray) and np.dtype == np.str_:
                            v = str(v)

                        while isinstance(v, basestring):
                            # strings are seperated by ';'
                            # XXX what if this is part of the value
                            # intended by the user?
                            v = v.split(';')

                        if expected_length != len(v):
                            raise ValueError("Unexpected length: %d != %d" %
                                                (expected_length, len(v)))

                        v = ArrayCollectable(v, length=expected_length)

                    collection[short_k] = v
                    continue

        found_label = False

        for label, collection in labels_collections:
            if k in label:
                collection[k] = v
                found_label = True
                break

        if found_label:
            continue

        # try to be smart and deduce this from dimensions.
        # this only works if nfeatures!=nsamples otherwise it would be
        # ambiguous
        # XXX is this ugly?
        if nfeatures != nsamples:
            try:
                n = len(v)
                if n == nfeatures:
                    fa[k] = v
                    continue
                elif n == nsamples:
                    sa[k] = v
                    continue
            except:
                pass

        # don't know what this is - make it a general attribute
        a[k] = v

    ds = Dataset(np.transpose(data), sa=sa, fa=fa, a=a)

    return ds

def to_niml(ds):
    '''Convert a Dataset to a NIML dataset

    Parameters
    ----------
    dataset: mvpa2.base.Dataset
        A PyMVPA Dataset

    Returns
    -------
    dset: dict
        Dictionary with NIML key-value pairs, such as obtained from
        mvpa2.support.nibabel.afni_niml_dset.read()
     '''
    if isinstance(ds, np.ndarray):
        ds = Dataset(ds)

    dset = dict(data=np.transpose(ds.samples))

    node_indices_labels = ('node_indices', 'center_ids', 'ids', 'roi_ids')
    node_indices = _find_node_indices(ds, node_indices_labels)
    if not node_indices is None:
        dset['node_indices'] = node_indices

    sample_labels = ('labels', 'targets')
    labels = _find_sample_labels(ds, sample_labels)
    if not labels is None:
        dset['labels'] = labels

    attr_labels = ('a', 'fa', 'sa')

    # a few labels are directly used in NIML dsets
    # without prefixing it with a pyMVPA string
    # for (dataset, feature, sample) attributes
    # here we define two for sample attributes
    attr_special_labels = ([], [], ['labels', 'stats'])

    for i, attr_label in enumerate(attr_labels):
        attr = getattr(ds, attr_label)
        special_labels = attr_special_labels[i]
        for k in attr.keys():
            v = attr[k]
            if hasattr(v, 'value'):
                v = v.value

            if k in special_labels:
                long_key = k
            else:
                long_key = _PYMVPA_SEP.join((_PYMVPA_PREFIX,
                                             attr_label.upper(), k))

            dset[long_key] = v

    return dset

def hstack(dsets, pad_to_feature_index=None, hstack_method='drop_nonunique',
                set_empty_value=0.):
    '''Stacks NIML datasets while considering node indices

    Parameters
    ----------
    dsets: list
        datasets to be stacked
    pad_to_feature_index: list or int or None
        If a list then it should be of the same length as dsets and indicates
        to which node index the input should be padded. A single int means
        that the same value is used for all dset in dsets. None means
        no padding, and is only allowed for non-sparse datasets.
    hstack_method: str:
        How datasets are stacked; see dataset.hstack.
    set_empty_value: float
        Value to which empty (padded) dataset values are set.

    Returns
    dset: Dataset
        Data combined from all dset in dsets.
    '''

    n = len(dsets)

    # make sure pad_to_feature_index has n values
    if pad_to_feature_index is None or type(pad_to_feature_index) is int:
        pad_to_feature_index = [pad_to_feature_index] * n
    elif len(pad_to_feature_index) != n:
        raise ValueError("illegal pad_to_feature_index: expected list or int")

    # labels that can contain node indices
    node_indices_labels = ('node_indices', 'center_ids', 'ids', 'roi_ids')
    node_indices = []

    # allocate space for output
    padded_dsets = []
    hstack_indices = []
    first_node_index = 0
    for i, (dset, pad_to) in enumerate(zip(dsets, pad_to_feature_index)):
        # get node indices in this dataset
        node_index = _find_node_indices(dset, node_indices_labels)
        if node_index is None:
            node_index = np.arange(dset.nfeatures)
        max_node_index = np.max(node_index)

        # make a stripped version - without node index labels
        stripped_dset = dset.copy()
        for label in node_indices_labels:
            if label in stripped_dset.fa:
                stripped_dset.fa.pop(label)

        # see if padding is needed
        if pad_to is None or pad_to == max_node_index + 1:
            if not np.array_equal(np.arange(max_node_index + 1), np.sort(node_index)):
                raise ValueError("Sparse input %d: need pad_to input" % (i + 1))
            padded_dset = stripped_dset
            other_index = np.arange(0)
        else:
            # have to use empty values
            nfeatures_empty = pad_to - dset.nfeatures
            if nfeatures_empty < 0:
                raise ValueError("Dataset has %d features, cannot pad "
                                    "to %d" % (dset.nfeatures, pad_to))

            # make empty array
            empty_arr = np.zeros((dset.nsamples, nfeatures_empty),
                                    dtype=dset.samples.dtype) + set_empty_value
            empty_dset = Dataset(empty_arr, sa=stripped_dset.sa.copy(deep=True))

            # combine current dset and empty array
            padded_dset = dataset.hstack((stripped_dset, empty_dset), hstack_method)

            # set the proper node indices
            other_index = np.setdiff1d(np.arange(pad_to), node_index)

        # sanity check to make sure that indices are ok
        # XXX could be more informative
        if len(np.setdiff1d(node_index, np.arange(pad_to or max_node_index + 1))):
            raise ValueError("Illegal indices")

        hstack_index = node_index + first_node_index
        hstack_other_index = other_index + first_node_index
        first_node_index += pad_to or (max_node_index + 1) # prepare for next iteration

        padded_dsets.append(padded_dset)
        hstack_indices.append(hstack_index)
        if len(other_index):
            hstack_indices.append(hstack_other_index)

    hstack_dset = dataset.hstack(padded_dsets, hstack_method)
    hstack_indices = np.hstack(hstack_indices)

    hstack_dset.fa[node_indices_labels[0]] = hstack_indices

    return hstack_dset


def _find_sample_labels(dset, sample_labels):
    '''Helper function to find labels in this dataset.
    Looks for any in sample_labels and returns the first one
    that matches '''
    use_label = None

    dset_keys = dset.sa.keys()
    for label in sample_labels:
        if label in dset_keys:
            sample_label = dset.sa[label].value
            if isinstance(sample_label, basestring):
                # split using
                sample_label = sample_label.split(';')

            # they can be of any type so ensure they are strings
            sample_label_list = [str(i) for i in sample_label]
            if len(sample_label_list) != dset.nsamples:
                # unlike node indices here we are more lenient
                # so not throw an exception but just continue
                continue

            use_label = label

            # do not look for any other labels
            break

    return None if use_label is None else sample_label_list


def _find_node_indices(dset, node_indices_labels):
    '''Helper function to find node indices in this dataset
    Sees if any of the node_indices_labels is a feature attribute
    in the dataset and returns it. If they are multiple matches
    ensure they are identical, otherwise raise an error.
    A use case is searchlight results that assignes center_ids as
    a feature attributes, but it should be named node_indices
    before conversion to NIML format'''

    use_label = None

    dset_keys = dset.fa.keys()
    for label in node_indices_labels:
        if label in dset_keys:
            if use_label is None:
                # make vector and ensure all integer values
                node_indices = dset.fa[label].value
                node_indices = np.asarray(node_indices).ravel()
                if len(node_indices) != dset.nfeatures:
                    raise ValueError("Node indices mismatch: found %d values "
                                     " but dataset has %d features" %
                                     (len(node_indices, dset.nfeatures)))
                node_indices_int = np.asarray(node_indices, dtype=np.int)
                if not np.array_equal(node_indices_int, node_indices):
                    raise ValueError("Node indices should have integer values")
                use_label = label

            else:
                if not np.array_equal(dset.fa[label].value, node_indices_int):
                    raise ValueError("Different indices for feature attributes"
                                     " %s and %s" % (use_key, label))

    return None if use_label is None else node_indices_int

def write(fn, ds, form='binary'):
    '''Write a Dataset to a file in NIML format

    Parameters
    ----------
    fn: str
        Filename
    ds: mvpa2.base.Dataset
        Dataset to be stored
    form: str
        Data format: 'binary' or 'text' or 'base64'
    '''
    niml_ds = to_niml(ds)
    niml_dset.write(fn, niml_ds, form=form)

def read(fn):
    '''Read a Dataset from a file in NIML format

    Parameters
    ----------
    fn: str
        Filename
    '''

    readers_converters = [(niml_dset.read, from_niml)]
    if externals.exists('h5py'):
        readers_converters.append((h5load, None))

    for reader, converter in readers_converters:
        try:
            r = reader(fn)
            if converter:
                r = converter(r)
            return r

        except:
            pass

    raise ValueError("Unable to read %s" % fn)


def from_any(x):
    '''Get a Dataset from the input

    Parameters
    ----------
    x: str or dict or Dataset
        Filename, or NIML-dictionary, or a Dataset itself

    Returns
    -------
    ds: mvpa2.base.Dataset
        Dataset instance
    '''
    if isinstance(x, basestring):
        return read(x)
    elif isinstance(x, dict):
        return from_niml(x)
    elif isinstance(x, Dataset):
        return x

    raise ValueError("Not supported: %r" % (x,))

