text: # 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.
  #
  ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
  """Support for storage using the NeuroImaging Markup Language (NIML).
  
  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 _as_vector_if_matrix_with_single_column(x):
      """Helper function"""
      if isinstance(x, np.ndarray) and len(x.shape) == 2 and x.shape[1] == 1:
          return x.ravel()
      else:
          return x
  
  
  def from_niml(dset, fa_labels=None, sa_labels=None, a_labels=None):
      """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
      """
      if fa_labels is None:
          fa_labels = []
      if sa_labels is None:
          sa_labels = []
      if a_labels is None:
          a_labels = []
  
      # 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.items():
          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 collection is not 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, str):
                              # 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 = _as_vector_if_matrix_with_single_column(v)
  
                          v = ArrayCollectable(v, length=expected_length)
  
                      collection[short_k] = v
                      continue
  
          v = _as_vector_if_matrix_with_single_column(v)
  
          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 node_indices is not None:
          dset["node_indices"] = node_indices
  
      sample_labels = ("labels", "targets")
      labels = _find_sample_labels(ds, sample_labels)
      if labels is not 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 list(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.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 = list(dset.sa.keys())
      for label in sample_labels:
          if label in dset_keys:
              sample_label = dset.sa[label].value
              if isinstance(sample_label, str):
                  # 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 = list(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_label, 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 = {(".dset",): (niml_dset.read, from_niml)}
  
      if externals.exists("h5py"):
          readers_converters[(".h5py", ".hdf")] = (h5load, None)
  
      keys = [
          exts
          for exts in readers_converters.keys()
          if any(fn.endswith(ext) for ext in exts)
      ]
  
      n_keys = len(keys)
  
      if n_keys == 1:
          # single extension matches
          key = keys[0]
          reader, converter = readers_converters[key]
  
          r = reader(fn)
          if converter is not None:
              r = converter(r)
          return r
  
      else:
          # unclear extension, try all readers and throw less informative
          # error message
          for reader, converter in readers_converters.values():
              try:
                  r = reader(fn)
                  if converter:
                      r = converter(r)
                  return r
  
              except:
                  pass
  
          raise ValueError("Unable to read %s with unclear extension" % 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, str):
          return read(x)
      elif isinstance(x, dict):
          return from_niml(x)
      elif isinstance(x, Dataset):
          return x
  
      raise ValueError("Not supported: %r" % (x,))
  