# 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.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
"""Base classes for Anatomy atlases support

TODOs:
======

 * major optimization. Now code is sloppy and slow -- plenty of checks etc

Module Organization
===================
mvpa2.atlases.base module contains support for various atlases

:group Base: BaseAtlas XMLBasedAtlas Label Level LabelsLevel
:group Talairach: PyMVPAAtlas LabelsAtlas ReferencesAtlas
:group Exceptions: XMLAtlasException

"""

import os.path as osp
from mvpa2.base import externals

if externals.exists('lxml', raise_=True, exception=ImportError):
    from lxml import etree, objectify

from mvpa2.base.dochelpers import enhanced_doc_string

import re
import numpy as np
from numpy.linalg import norm

from mvpa2.atlases.transformation import SpaceTransformation, Linear
from mvpa2.misc.support import reuse_absolute_path

if externals.exists('nibabel', raise_=True):
    import nibabel as nb

from mvpa2.base import warning
if __debug__:
    from mvpa2.base import debug


def check_range(coord, range):
    """Check if coordinates are within range (0,0,0) - (range)

    Returns
    -------
    bool
      Success status
    """
    # TODO: optimize
    if len(coord) != len(range):
        raise ValueError("Provided coordinate %r and given range %r" % \
                         (coord, range) + \
                         " have different dimensionality"
                         )
    for c, r in zip(coord, range):
        if c < 0 or c >= r:
            return False
    return True

#
# Base classes
#

class XMLAtlasException(Exception):
    """Exception to be thrown if smth goes wrong dealing with XML based atlas
    """
    pass

class BaseAtlas(object):
    """Base class for the atlases.
    """
    pass



class XMLBasedAtlas(BaseAtlas):
    """Base class for atlases using XML as the definition backend
    """

    def __init__(self, filename=None,
                 resolution=None, image_file=None,
                 coordT=None, default_levels=None, load_maps=True):
        """
        Parameters
        ----------
        filename : str
          Filename for the xml definition of the atlas
        resolution : None or float
          Some atlases link to multiple images at different
          resolutions. if None -- best resolution is selected
          using 0th dimension resolution
        image_file : None or str
          If None, overrides filename for the used imagefile, so
          it could load a custom (re-registered) atlas maps
        coordT
          Optional transformation to apply first
        default_levels : None or slice or list of int
          What levels by default to operate on
        load_maps : bool
          Load spatial maps for the atlas.
        """
        BaseAtlas.__init__(self)

        self.__atlas = None

        self._image_file = None
        self._filename = filename
        # TODO: think about more generalizable way?
        self._resolution = resolution
        self._force_image_file = image_file
        self.default_levels = default_levels

        if filename:
            self.load_atlas(filename)

        # common sanity checks
        if not self._check_version(self.version):
            raise IOError(
                "Version %s is not recognized to be native to class %s" % \
                (self.version, self.__name__))

        if not set(['header', 'data']) \
               == set([i.tag for i in self.getchildren()]):
            raise IOError("No header or data were defined in %s" % filename)

        header = self.header
        headerChildrenTags = XMLBasedAtlas._children_tags(header)
        if not ('images' in headerChildrenTags) or \
           not ('imagefile' in XMLBasedAtlas._children_tags(header.images)):
            raise XMLAtlasException(
                "Atlas requires image/imagefile header fields")

        # Load and post-process images
        self._image = None
        if load_maps:
            self._load_images()
            if self._image is not None:
                # Get extent and voxel dimensions, limiting to 3D
                self._extent = np.abs(np.asanyarray(self._image.shape[:3]))
                self._voxdim = np.asanyarray(self._image.header.get_zooms()[:3])
                self.relativeToOrigin = True
        # Assign transformation to get into voxel coordinates,
        # spaceT will be set accordingly
        self.set_coordT(coordT)
        self._load_metadata()


    def _check_range(self, c):
        """ check and adjust the voxel coordinates"""
        # check range
        if __debug__:
            debug('ATL__', "Querying for voxel %r" % (c,))
        if not check_range(c, self.extent):
            warning("Coordinates %r are not within the extent %r." \
                    " Reseting to (0,0,0)" % (c, self.extent))
            # assume that voxel [0,0,0] is blank, i.e. carries
            # no labels which could possibly result in evil outcome
            c = [0]*3
        return c


    @staticmethod
    def _check_version(version):
        """To be overriden in the derived classes. By default anything is good"""
        return True

    def _load_images(self):
        """To be overriden in the derived classes. By default does nothing"""
        pass

    def _load_metadata(self):
        """To be overriden in the derived classes. By default does nothing"""
        pass

    def load_atlas(self, filename):
        """Load atlas from a file
        """
        if __debug__:
            debug('ATL_', "Loading atlas definition xml file " + filename)
        # Create objectify parser first
        parser = etree.XMLParser(remove_blank_text=True)
        lookup = objectify.ObjectifyElementClassLookup()
        parser.setElementClassLookup(lookup)
        try:
            self.__atlas = etree.parse(filename, parser).getroot()
        except IOError:
            raise XMLAtlasException("Failed to load XML file %s" % filename)

    @property
    def version(self):
        if self.__atlas is not None \
               and ("version" in list(self.__atlas.attrib.keys())):
            return self.__atlas.get("version")
        else:
            return None

    @staticmethod
    def _children_tags(root):
        """Little helper to return tags for the children of the node
        """
        return [i.tag for i in root.getchildren()]


    def __getattr__(self, attr):
        """
        Lazy way to provide access to the definitions in the atlas
        """
        if self.__atlas is not None:
            return getattr(self.__atlas, attr)
        else:
            raise XMLAtlasException(
                "Atlas in " + self.__name__ + " was not read yet")


    def set_coordT(self, coordT):
        """Set coordT transformation.

        spaceT needs to be adjusted since we glob those two
        transformations together
        """
        self._coordT = coordT           # lets store for debugging etc
        if self._image is not None:
            # Combine with the image's qform
            coordT = Linear(np.linalg.inv(self._image.header.get_qform()),
                            previous=coordT)
        self._spaceT = SpaceTransformation(
            previous=coordT, to_real_space=False)


    ##REF: Name was automagically refactored
    def label_point(self, coord, levels=None):
        """Return labels for the given spatial point at specified levels

        Function takes care about first transforming the point into
        the voxel space

        Parameters
        ----------
        coord : tuple
          Coordinates of the point (xyz)
        levels : None or list of int
          At what levels to return the results
        """
        coord_ = np.asarray(coord)          # or we would alter what should be constant
        #if not isinstance(coord, np.numpy):
        #c = self.getVolumeCoordinate(coord)
        #c = self.spaceT.to_voxel_space(coord_)
        #if self.coordT:
        #   coord_t = self.coordT[coord_]
        #else:
        #   coord_t = coord_

        c = self.spaceT(coord_)

        result = self.label_voxel(c, levels)
        result['coord_queried'] = coord
        #result['coord_trans'] = coord_t
        result['voxel_atlas'] = c
        return result


    def levels_listing(self):
        lkeys = list(range(self.nlevels))
        return '\n'.join(['%d: ' % k + str(self._levels[k])
                          for k in lkeys])


    def _get_selected_levels(self, levels=None):
        """Helper to provide list of levels to operate on

        Depends on given `levels` as well as self.default_levels
        """
        if levels is None:
            levels = [ i for i in range(self.nlevels) ]
        elif (isinstance(levels, slice)):
            # levels are given as a range
            if levels.step: step = levels.step
            else: step = 1

            if levels.start: start = levels.start
            else: start = 0

            if levels.stop: stop = levels.stop
            else: stop = self.nlevels

            levels = [ i for i in range(start, stop, step) ]

        elif isinstance(levels, (list, tuple)):
            # levels given as list
            levels = list(levels)

        elif isinstance(levels, int):
            levels = [ levels ]

        else:
            raise TypeError('Given levels "%s" are of unsupported type' % repr(levels))

        selected_levels = levels
        # test given values
        levels = self.levels
        for level in selected_levels:
            if not level in levels:
                raise ValueError("Level %r is not known (out of range?). Known levels are:\n%s" \
                      % (level, self.levels_listing()))

        return selected_levels


    def query(self, index, query_voxel=False):
        """Generic query method.

        Use shortcuts `__getitem__` for querying by voxel indices and
        `__call__` for querying by space coordinates.

        Parameters
        ----------
        index : tuple or list
          Arguments of the query, such as coordinates and optionally
          levels
        query_voxel : bool
          Either at the end query a voxel indexes or point coordinates

        Allows to access the elements via simple indexing. Examples::

            print atlas[ 0, -7, 20, [1,2,3] ]
            print atlas[ (0, -7, 20), 1:2 ]
            print atlas[ (0, -7, 20) ]
            print atlas[ (0, -7, 20), : ]
        """
        if len(index) in [2, 4]:
            levels_slice = index[-1]
        else:
            if self.default_levels is None:
                levels_slice = slice(None, None, None)
            else:
                levels_slice = self.default_levels

        levels = self._get_selected_levels(levels=levels_slice)

        if len(index) in [3, 4]:
            # we got coordinates 1 by 1 + may be a level
            coord = index[0:3]

        elif len(index) in [1, 2]:
            coord = index[0]
            if isinstance(coord, (list, tuple)):
                if len(coord) != 3:
                    raise TypeError("Given coordinates must be in 3D")
            else:
                raise TypeError("Given coordinates must be a list or a tuple")

        else:
            raise TypeError("Unknown shape of parameters `%s`" % repr(index))

        if query_voxel:
            return self.label_voxel(coord, levels)
        else:
            return self.label_point(coord, levels)

    #
    ## Shortcuts for `query`
    #
    def __getitem__(self, index):
        """Query atlas with voxel indexes

        Examples
        --------

        ::
            print atlas[ 0, -7, 20, [1,2,3] ]
            print atlas[ (0, -7, 20), 1:2 ]
            print atlas[ (0, -7, 20) ]
            print atlas[ (0, -7, 20), : ]
        """
        return self.query(index, True)

    def __call__(self, *args):
        return self.query(args, False)

    # REDO in some sane fashion so referenceatlas returns levels for the base
    def _get_levels(self):
        return self._get_levels_virtual()

    ##REF: Name was automagically refactored
    def _get_levels_virtual(self):
        return self._levels

    levels = property(fget=_get_levels)


    resolution = property(fget=lambda self:self._resolution)
    origin = property(fget=lambda self:self._origin)
    extent = property(fget=lambda self:self._extent)
    voxdim = property(fget=lambda self:self._voxdim)
    spaceT = property(fget=lambda self:self._spaceT)
    coordT = property(fget=lambda self:self._spaceT,
                      fset=set_coordT)

class Label(object):
    """Represents a label. Just to bring all relevant information together
    """
    def __init__ (self, text, abbr=None, coord=(None, None, None),
                  count=0, index=0):
        """
        Parameters
        ----------
        text : str
          Fullname of the label
        abbr : str, optional
          Abbreviated name.
        coord : tuple of float, optional
          Coordinates.
        count : int, optional
          Count of those labels in the atlas
        """
        self.text = text.strip()
        if abbr is not None:
            abbr = abbr.strip()
        self.coord = coord
        self.count = count
        self.__abbr = abbr
        self.__index = int(index)


    @property
    def index(self):
        return self.__index

    def __repr__(self):
        return "Label(%r%s, coord=%r, count=%r, index=%r)" % \
               (self.text,
                (', abbr=%s' % repr(self.__abbr), '')[int(self.__abbr is None)],
                self.coord, self.count, self.__index)

    def __str__(self):
        return self.text

    @staticmethod
    def from_xml(Elabel):
        """Create label from an XML node
        """
        kwargs = {}
        if 'x' in Elabel.attrib:
            kwargs['coord'] = ( Elabel.attrib.get('x'),
                                Elabel.attrib.get('y'),
                                Elabel.attrib.get('z') )
        for l in ('count', 'abbr', 'index'):
            if l in Elabel.attrib:
                kwargs[l] = Elabel.attrib.get(l)
        return Label(Elabel.text.strip(), **kwargs)

    @property
    def abbr(self):
        """Returns abbreviated version if such is available
        """
        if self.__abbr in [None, ""]:
            return self.text
        else:
            return self.__abbr


class Level(object):
    """Represents a level. Just to bring all relevant information together
    """
    def __init__ (self, description):
        self.description = description
        self._type = "Base"

    def __repr__(self):
        return "%s Level: %s" % \
               (self.level_type, self.description)

    def __str__(self):
        return self.description

    @staticmethod
    ##REF: Name was automagically refactored
    def from_xml(Elevel, level_type=None):
        """Simple factory of levels
        """
        if level_type is None:
            if not 'type' in Elevel.attrib:
                raise XMLAtlasException("Level must have type specified. Level: " + repr(Elevel))
            level_type = Elevel.get("type")

        levelTypes = { 'label':     LabelsLevel,
                       'reference': ReferencesLevel }

        if level_type in levelTypes:
            return levelTypes[level_type].from_xml(Elevel)
        else:
            raise XMLAtlasException("Unknown level type " + level_type)

    level_type = property(lambda self: self._type)


class LabelsLevel(Level):
    """Level of labels.

	XXX extend
    """
    def __init__ (self, description, index=None, labels=None):
        if labels is None:
            labels = []
        Level.__init__(self, description)
        self.__index = index
        self.__labels = labels
        self._type = "Labels"

    def __repr__(self):
        return Level.__repr__(self) + " [%d] " % \
               (self.__index)

    @staticmethod
    ##REF: Name was automagically refactored
    def from_xml(Elevel, levelIndex=None):
        # XXX this is just for label type of level. For distance we need to ...
        # we need to assure the right indexing

        if levelIndex is None:
                levelIndex = [0]
        index = 0
        if 'index' in Elevel.attrib:
            index = int(Elevel.get("index"))

        maxindex = max([int(i.get('index')) \
                        for i in Elevel.label[:]])
        labels = [ None for i in range(maxindex+1) ]
        for label in Elevel.label[:]:
            labels[ int(label.get('index')) ] = Label.from_xml(label)

        levelIndex[0] = max(levelIndex[0], index) + 1 # assign next one

        return LabelsLevel(Elevel.get('description'),
                           index,
                           labels)

    @property
    def index(self): return self.__index

    @property
    def labels(self): return self.__labels

    def __getitem__(self, index):
        return self.__labels[index]

    def find(self, target, unique=True):
        """Return labels descr of which matches the string

        Parameters
        ----------
        target : str or re._pattern_type
          Substring in abbreviation to be searched for, or compiled
          regular expression to be searched or matched if anchored.
        unique : bool, optional
          If True, raise exception if none or more than 1 was found. Return
          just a single item if found (not list).
        """
        if isinstance(target, re._pattern_type):
            res = [l for l in self.__labels if target.search(l.abbr)]
        else:
            res = [l for l in self.__labels if target in l.abbr]

        if unique:
            if len(res) != 1:
                raise ValueError("Got %d matches whenever just 1 was " \
                      "looked for (target was %s)." % (len(res), target))
            return res[0]
        else:
            return res


class ReferencesLevel(Level):
    """Level which carries reference points
    """
    def __init__ (self, description, indexes=None):
        if indexes is None:
            indexes = []
        Level.__init__(self, description)
        self.__indexes = indexes
        self._type = "References"

    @staticmethod
    ##REF: Name was automagically refactored
    def from_xml(Elevel):
        # XXX should probably do the same for the others?
        requiredAttrs = ['x', 'y', 'z', 'type', 'description']
        if not set(requiredAttrs) == set(Elevel.attrib.keys()):
            raise XMLAtlasException("ReferencesLevel has to have " +
                                    "following attributes defined " +
                                    repr(requiredAttrs))

        indexes = tuple(int(Elevel.get(a)) for a in ('x', 'y', 'z'))

        return ReferencesLevel(Elevel.get('description'),
                               indexes)

    @property
    def indexes(self):
        return self.__indexes


class PyMVPAAtlas(XMLBasedAtlas):
    """Base class for PyMVPA atlases, such as LabelsAtlas and ReferenceAtlas
    """

    source = 'PyMVPA'

    def __init__(self, *args, **kwargs):
        XMLBasedAtlas.__init__(self, *args, **kwargs)

        # sanity checks
        header = self.header
        headerChildrenTags = XMLBasedAtlas._children_tags(header)
        if not ('space' in headerChildrenTags) or \
           not ('space-flavor' in headerChildrenTags):
            raise XMLAtlasException("PyMVPA Atlas requires specification of" +
                                    " the space in which atlas resides")

        self.__space = header.space.text
        self.__spaceFlavor = header['space-flavor'].text


    __doc__ = enhanced_doc_string('PyMVPAAtlas', locals(), XMLBasedAtlas)


    ##REF: Name was automagically refactored
    def _load_images(self):
        # shortcut
        imagefile = self.header.images.imagefile
        #self.nlevels = len(self._levels_by_id)

        # Set offset if defined in XML file
        # XXX: should just take one from the qoffset... now that one is
        #       defined... this origin might be misleading actually
        self._origin = np.array( (0, 0, 0) )
        if 'offset' in imagefile.attrib:
            self._origin = np.array( [int(x) for x in
                                     imagefile.get('offset').split(',')] )

        # Load the image file which has labels
        if self._force_image_file is not None:
            imagefilename = self._force_image_file
        else:
            imagefilename = imagefile.text
        imagefilename = reuse_absolute_path(self._filename, imagefilename)

        try:
            self._image = None
            for ext in ['', '.nii.gz']:
                try:
                    self._image  = nb.load(imagefilename + ext)
                    break
                except Exception as e:
                    pass
            if self._image is None:
                raise e
        except RuntimeError as e:
            raise RuntimeError(" Cannot open file %s due to %s" % (imagefilename, e))

        self._data = self._image.get_data()
        # we get the data as x,y,z[,t] but we want to have the time axis first
        # if any
        if len(self._data.shape) == 4:
            self._data = np.rollaxis(self._data, -1)

        # remove bogus dimensions on top of 4th
        if len(self._data.shape[0:-4]) > 0:
            bogus_dims = self._data.shape[0:-4]
            if max(bogus_dims)>1:
                raise RuntimeError("Atlas %s has more than 4 of non-singular" \
                      "dimensions" % imagefilename)
            new_shape = self._data.shape[-4:]
            self._data.reshape(new_shape)

        #if self._image.extent[3] != self.nlevels:
        #   raise XMLAtlasException("Atlas %s has %d levels defined whenever %s has %d volumes" % \
        #                           ( filename, self.nlevels, imagefilename, self._image.extent[3] ))


    ##REF: Name was automagically refactored
    def _load_metadata(self):
        # Load levels
        self._levels = {}
        # preprocess labels for different levels
        self._Nlevels = 0
        index_incr = 0
        for index, child in enumerate(self.data.getchildren()):
            if child.tag == 'level':
                level = Level.from_xml(child)
                self._levels[level.description] = level
                if hasattr(level, 'index'):
                    index = level.index
                else:
                    # to avoid collision if some levels do
                    # have indexes
                    while index_incr in self._levels:
                        index_incr += 1
                    index, index_incr = index_incr, index_incr+1
                self._levels[index] = level
            else:
                raise XMLAtlasException(
                    "Unknown child '%s' within data" % child.tag)
            self._Nlevels += 1


    ##REF: Name was automagically refactored
    def _get_nlevels_virtual(self):
        return self._Nlevels

    ##REF: Name was automagically refactored
    def _get_nlevels(self):
        return self._get_nlevels_virtual()

    @staticmethod
    ##REF: Name was automagically refactored
    def _check_version(version):
        # For compatibility lets support "RUMBA" atlases
        return version.startswith("pymvpa-") or version.startswith("rumba-")


    space = property(fget=lambda self:self.__space)
    space_flavor = property(fget=lambda self:self.__spaceFlavor)
    nlevels = property(fget=_get_nlevels)


class LabelsAtlas(PyMVPAAtlas):
    """
    Atlas which provides labels for the given coordinate
    """

    ##REF: Name was automagically refactored
    def label_voxel(self, c, levels=None):
        """
        Return labels for the given voxel at specified levels specified by index
        """
        levels = self._get_selected_levels(levels=levels)

        result = {'voxel_queried' : c}

        # check range
        c = self._check_range(c)

        resultLevels = []
        for level in levels:
            if level in self._levels:
                level_ = self._levels[ level ]
            else:
                raise IndexError(
                    "Unknown index or description for level %d" % level)

            resultIndex =  int(self._data[ level_.index, \
                                            c[0], c[1], c[2] ])

            resultLevels += [ {'index': level_.index,
                               'id': level_.description,
                               'label' : level_[ resultIndex ]} ]

        result['labels'] = resultLevels
        return result

    __doc__ = enhanced_doc_string('LabelsAtlas', locals(), PyMVPAAtlas)



class ReferencesAtlas(PyMVPAAtlas):
    """
    Atlas which provides references to the other atlases.

    Example: the atlas which has references to the closest points
    (closest Gray, etc) in another atlas.
    """

    def __init__(self, distance=0, reference_level=None, *args, **kwargs):
        """Initialize `ReferencesAtlas`
        """
        PyMVPAAtlas.__init__(self, *args, **kwargs)
        # sanity checks
        if not ('reference-atlas' in XMLBasedAtlas._children_tags(self.header)):
            raise XMLAtlasException(
                "ReferencesAtlas must refer to a some other atlas")

        referenceAtlasName = self.header["reference-atlas"].text

        # uff -- another evil import but we better use the factory method
        from mvpa2.atlases.warehouse import Atlas
        self.__referenceAtlas = Atlas(filename=reuse_absolute_path(
            self._filename, referenceAtlasName))

        if self.__referenceAtlas.space != self.space or \
           self.__referenceAtlas.space_flavor != self.space_flavor:
            raise XMLAtlasException(
                "Reference and original atlases should be in the same space")

        self.__referenceLevel = None    # pylint shut up
        if reference_level is not None:
            self.set_reference_level(reference_level)
        self.set_distance(distance)

    __doc__ = enhanced_doc_string('ReferencesAtlas', locals(), PyMVPAAtlas)

    # number of levels must be of the referenced atlas due to
    # handling of that in __getitem__
    #nlevels = property(fget=lambda self:self.__referenceAtlas.nlevels)
    ##REF: Name was automagically refactored
    def _get_nlevels_virtual(self):
        return self.__referenceAtlas.nlevels


    ##REF: Name was automagically refactored
    def set_reference_level(self, level):
        """
        Set the level which will be queried
        """
        if level in self._levels:
            self.__referenceLevel = self._levels[level]
        else:
            raise IndexError("Unknown reference level %r. " % level + \
                  "Known are %r" % (list(self._levels.keys()), ))


    ##REF: Name was automagically refactored
    def label_voxel(self, c, levels = None):

        if self.__referenceLevel is None:
            warning("You did not provide what level to use "
                    "for reference. Assigning 0th level -- '%s'"
                    % (self._levels[0],))
            self.set_reference_level(0)
            # return self.__referenceAtlas.label_voxel(c, levels)

        c = self._check_range(c)

        # obtain coordinates of the closest voxel
        cref = self._data[ self.__referenceLevel.indexes, c[0], c[1], c[2] ]
        dist = norm( (cref - c) * self.voxdim )
        if __debug__:
            debug('ATL__', "Closest referenced point for %r is "
                  "%r at distance %3.2f" % (c, cref, dist))
        if (self.distance - dist) >= 1e-3: # neglect everything smaller
            result = self.__referenceAtlas.label_voxel(cref, levels)
            result['voxel_referenced'] = c
            result['distance'] = dist
        else:
            result = self.__referenceAtlas.label_voxel(c, levels)
            if __debug__:
                debug('ATL__', "Closest referenced point is "
                      "further than desired distance %.2f" % self.distance)
            result['voxel_referenced'] = None
            result['distance'] = 0
        return result


    ##REF: Name was automagically refactored
    def levels_listing(self):
        return self.__referenceAtlas.levels_listing()

    ##REF: Name was automagically refactored
    def _get_levels_virtual(self):
        return self.__referenceAtlas.levels

    ##REF: Name was automagically refactored
    def set_distance(self, distance):
        """Set desired maximal distance for the reference
        """
        if distance < 0:
            raise ValueError("Distance should not be negative. "
                             " Thus '%f' is not a legal value" % distance)
        if __debug__:
            debug('ATL__',
                  "Setting maximal distance for queries to be %d" % distance)
        self.__distance = distance

    distance = property(fget=lambda self:self.__distance, fset=set_distance)
    reference_level = property(fget=lambda self:self.__referenceLevel, fset=set_reference_level)

