#!/usr/bin/env python
# encoding: utf-8

'''
Several tools for look-up tables management and interpolation

Provides:
    - LUT class: Look-Up Tables
      extends ndarrays for generic look-up table management including dimensions description,
      binary operations, multi-dimensional interpolation, plotting, etc
    - Idx class: find the index of values, for LUT interpolation
    - MLUT class: Multiple Look-Up Tables
      groups several LUTs, provides I/O interfaces with multiple data types
    - merge: MLUT merging
'''

from __future__ import print_function, division, absolute_import
import sys
import numpy as np
from scipy.interpolate import interp1d
import xarray as xr
from os.path import exists
from os import remove
from collections import OrderedDict
from numpy.ma import filled
import warnings
import itertools
if sys.version_info[:2] >= (3, 0): # python2/3 compatibility
    unicode = str
    xrange = range


def interleave_seq(p, q):
    '''
    Interleave 2 sequences (union, preserve order)
    ([1, 3, 4, 6], [2, 3, 6]) -> [1, 2, 3, 4, 6]
    '''
    if len(p) == 0:
        return q
    elif len(q) == 0:
        return p
    elif p[0] == q[0]:
        return [p[0]] + interleave_seq(p[1:], q[1:])
    elif p[0] in q[1:]:
        if q[0] in p[1:]:
            raise ValueError('sequences "{}" and "{}" cannot be interleaved'.format(p, q))
        return interleave_seq(q, p)
    else:  # p[0] not in q
        return [p[0]] + interleave_seq(p[1:], q)

def sizeof_fmt(num, suffix='B'):
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)

def uniq(seq):
    '''
    Returns uniques elements from a sequence, whist preserving its order
    http://stackoverflow.com/questions/480214/
    '''
    seen = set()
    seen_add = seen.add
    return [ x for x in seq if not (x in seen or seen_add(x))]


def bin_edges(x, min=None, max=None):
    '''
    calculate n+1 bin edges from n bin centers in x
    '''
    assert x.ndim == 1
    if len(x) == 1:
        ret = np.array([x-0.5, x+0.5])
    else:
        first = (3*x[0] - x[1])/2.
        last = (3*x[-1] - x[-2])/2.
        ret = np.append(np.append(first, 0.5*(x[1:]+x[:-1])), last)
    
    if min is not None:
        ret = np.maximum(ret, min)
    if max is not None:
        ret = np.minimum(ret, max)
    
    return ret


class LUT(object):
    '''
    Look-up table storage with generic multi-dimensional interpolation.
    Extends the __getitem__ method of ndarrays to float and float arrays (index
    tables with floats)
    The LUT axes can be optionally provided so that values can be interpolated
    into float indices in a first step, using the Idx class.
    Other features:
        * binary operations between LUTs and with scalars
        * automatic broadcasting:
          binary operations between LUTs with different axes
          the LUTs are broadcasted together according to their common axes
        * apply functions, dimension reduction (apply a function along a given axis)
        * equality testing
        * plotting

    Arguments:
        * data is a n-dimension array containing the LUT data
        * axes is a list representing each dimension, containing for each
          dimension:
            - a 1-d array or list containing the tabulated values of the dimension
            - or None, if there are no values associated with this dimension
        * names is an optional list of names for each axis (default: None)
          it is necessary for storing the LUT
        * attr is a dictionary of additional attributes, useful for merging LUTs
        * desc is a string describing the parameter stored (default None)
          when using save(), desc is used as the hdf dataset name to store the LUT

    Attributes: axes, shape, data, ndim, attrs, names, desc

    Example 1
    ---------
    Basic usage, without axes and in 1D:
    >>> data = np.arange(10)**2
    >>> L = LUT(data)
    >>> L[1], L[-1], L[::2]    # standard indexing is possible
    (1, 81, array([ 0,  4, 16, 36, 64]))
    >>> L[1.5]    # Indexing with a float: interpolation
    2.5
    >>> L[np.array([[0.5, 1.5], [2.5, 9.]])]    # interpolate several values at once
    array([[  0.5,   2.5],
           [  6.5,  81. ]])

    Example 2
    ---------
    Interpolation of the atmospheric pressure
    >>> z = np.linspace(0, 120., 80)
    >>> P0 = np.linspace(980, 1030, 6)
    >>> Pdata = P0.reshape(1,-1)*np.exp(-z.reshape(-1,1)/8) # dimensions (z, P0)

    A 2D LUT with attached axes.  Axes names can optionally be provided.
    >>> P = LUT(Pdata, axes=[z, P0], names=['z', 'P0'])
    >>> P[Idx(8.848), Idx(1013.)]  # standard pressure at mount Everest
    336.09126751112842
    >>> z = np.random.rand(50, 50)  # now z is a 2D array of elevations between 0 and 1 km
    >>> P0 = 1000+20*np.random.rand(50, 50)  # and P0 is a 2D array of pressures at z=0
    >>> P[Idx(z), Idx(P0)].shape    # returns a 2D array of corresponding pressures
    (50, 50)

    In Idx, the values can optionally be passed using keyword notation.
    In this case, there is a verification that the argument corresponds to the right axis name.
    >>> P[:, Idx(1013., 'P0')].shape   # returns the (shape of) standard vertical pressure profile
    (80,)
    '''

    def __init__(self, data, axes=None, names=None, desc=None, attrs=None):
        self.data = data
        self.desc = desc
        if attrs is None:
            self.attrs = OrderedDict()
        else:
            self.attrs = attrs
        self.ndim = len(self.data.shape)
        self.shape = data.shape

        # check axes
        if axes is None:
            self.axes = self.ndim * [None]
        else:
            self.axes = axes
            assert len(axes) == self.ndim
            for i, ax in enumerate(axes):
                if isinstance(ax, np.ndarray):
                    assert ax.ndim == 1
                    assert len(ax) == data.shape[i]
                elif isinstance(ax, list):
                    assert len(ax) == data.shape[i]
                elif ax is None: pass
                else:
                    assert ax[:].ndim == 1
                    assert len(ax) == data.shape[i]

        # check names
        if names is None:
            self.names = self.ndim * [None]
        else:
            self.names = names
            assert len(names) == self.ndim

        if self.data.dtype in [float, np.float32, np.float64]:
            self.formatter = '{:3g}'
        else:
            self.formatter = '{}'

    def sub(self, d=None, ignore=False):
        '''
        returns a subset LUT of current LUT along several axes

        * d is a dictionary of {ax: value}, where:

            - ax is the name (string) or index (integer) of the axis to consider

            - value can be:
                * scalar integer (axis is removed)
                * float index (axis is removed with interpolation)
                * slice (axis is subsetted)
                * 1-d array (axis is subsetted)
                * Idx

        * ignore: if True, return full LUT if axis is not present in LUT

        Examples:
          lut.sub({'axis1': 1.5})
             returns the lut stripped from 'axis1', for which we use the index 1.5
          lut.sub({2: Idx(42.)})
          lut.sub({'ax': arange(5)})   # first 5 values
          lut.sub({'z': slice(None,None,2)})   # every other value
          lut.sub({'wav': lut.axis('wav')<500.})   # array of booleans
          lut.sub({'wav': Idx(lambda x: x<500.)})  # select wav < 500
        '''
        if d is None:
            return Subsetter(self)

        if self.ndim == 0:
            # scalar lut, cannot subset
            if ignore:
                return self
            else:
                raise Exception('Cannot subset scalar LUT {}'.format(self))

        keys = [slice(None)] * self.ndim
        names = list(self.names)
        axes = list(self.axes)
        dims_to_remove = []

        for ax, v in d.items():
            # ax: str or int
            if isinstance(ax, int):
                iax = ax
            else:
                assert isinstance(ax, str), 'ax should be str or int'
                if ax in self.names:
                    iax = self.names.index(ax)
                else:
                    if ignore:
                        continue
                    else:
                        raise Exception('sub is not possible on axis "{}" '
                                        'because this axis is not present '
                                        'in {}'.format(ax, self))

            if isinstance(v, Idx_base):
                idx = v.index(self.axes[iax])
                newax = v.apply(self.axes[iax])
            else:
                idx = v
                newax = None

            # axes: names and values
            if (np.array(idx).ndim == 0) and not isinstance(v, slice):
                # scalar, remove dimension
                dims_to_remove.append(iax)
            # in all other cases, keep dimension
            elif axes[iax] is None:
                axes[iax] = None
            elif isinstance(v, slice):
                axes[iax] = axes[iax][v]
            elif np.array(idx).ndim == 1:
                if newax is not None:
                    axes[iax] = newax
                elif v.dtype.kind in ['u', 'i']: # integer
                    axes[iax] = axes[iax][v]
                else:
                    axes[iax] = None
            else:
                raise Exception('Cannot use a {}-dim array to '
                                'subset a LUT'.format(np.array(idx).ndim))

            keys[iax] = idx

        axes = [a for i, a in enumerate(axes) if not i in dims_to_remove]
        names = [a for i, a in enumerate(names) if not i in dims_to_remove]

        data = self[tuple(keys)]

        return LUT(data, axes=axes, names=names,
                   attrs=dict(self.attrs), desc=self.desc)


    def axis(self, a, aslut=False):
        '''
        returns axis referred to by a (string or integer)
        aslut:
            False: returns the values
            True: returns axis a a LUT
                (containing itself as only axis)
        '''
        if isinstance(a, str):
            index = self.names.index(a)
        elif isinstance(a, int):
            index = a
        else:
            raise TypeError('argument of LUT.axis() should be int or string')

        if aslut:
            data = self.axes[index]
            return LUT(data, axes=[data], names=[self.names[index]])
        else:
            return self.axes[index]

    def print_info(self, *args, **kwargs):
        # same as describe()
        return self.describe(*args, **kwargs)

    def describe(self, show_attrs=False):
        '''
        Prints the LUT informations
        arguments:
            show_attrs: show LUT attributes (default: False)

        returns self
        '''
        try:
            rng = ' between {:.3g} and {:.3g}'.format(
                    np.amin(self.data), np.amax(self.data)
                    )
        except:
            rng = ''

        print('LUT {}({}{}):'.format(
                {True: '', False: '"{}" '.format(self.desc)}[self.desc is None],
                self.data.dtype, rng,
                ))

        for i in xrange(len(self.data.shape)):
            if self.names[i] is None:
                name = 'NoName'
            else:
                name = self.names[i]
            if self.axes[i] is None:
                print('  Dim {} ({}): {} values, no axis attached'.format(i, name, self.data.shape[i]))
            else:
                print('  Dim {} ({}): {} values in [{}, {}]'.format(
                        i, name,
                        len(self.axes[i]),
                        self.axes[i][0],
                        self.axes[i][-1]))
        if show_attrs:
            print(' Attributes:')
            for k, v in self.attrs.items():
                print(' ', k, ':', v)

        return self

    def __getitem__(self, keys):
        '''
        Get items from the LUT, with possible interpolation

        Indexing works mostly like a standard array indexing, with the following differences:
            - float indexes can be provided, they will result in an
              interpolation between bracketing integer indices for this dimension
            - float arrays can be provided, they will result in an
              interpolation between bracketing integer arrays for this dimension
            - Idx objects can be provided, they are first converted into indices
            - basic indexing and slicing, and advanced indexing (indexing with
              an ndarray of int or bool) can still be used
              (see http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html)
            - if arrays are passed as arguments, they must all have the same shape
            - Unlike ndarrays, the number of dimensions in keys should be
              identical to the dimension of the LUT
              >>> LUT(np.zeros((2, 2)))[:]
              Traceback (most recent call last):
              ...
              Exception: Incorrect number of dimensions in __getitem__

        Returns: a scalar or ndarray
        '''

        if self.data.dtype.char in ['S', 'U']:
            # string or unicode arrays: bypass this method and directly apply
            # ndarray.__getitem__
            return self.data[keys]

        if not isinstance(keys, tuple):
            keys = (keys,)
        keys = list(keys)
        N = len(keys)

        if N != self.ndim:
            raise Exception('Incorrect number of dimensions in __getitem__ '
                            '(expecting {}, got {})'.format(self.ndim, N))

        # convert the Idx keys to float indices
        for i in xrange(N):
            k = keys[i]
            if isinstance(k, Idx_base):
                # k is an Idx instance
                # convert it to float indices for the current axis
                if k.name not in [None, self.names[i]]:
                    msg = 'Error, wrong parameter passed at position {}, expected {}, got {}'
                    raise Exception(msg.format(i, self.names[i], k.name))
                keys[i] = k.index(self.axes[i])

        # determine the dimensions of the result (for broadcasting coef)
        dims_array = None
        index0 = []
        for i in xrange(N):
            k = keys[i]
            if isinstance(k, np.ndarray) and (k.ndim > 0):
                index0.append(np.zeros_like(k, dtype='int'))
                if dims_array is None:
                    dims_array = k.shape
                else:
                    assert dims_array == k.shape, 'LUTS.__getitem__: all arrays must have same shape ({} != {})'.format(str(dims_array), str(k.shape))
            elif isinstance(k, slice):
                index0.append(k)
            else:  # scalar
                index0.append(0)

        shp_res = np.zeros(1).reshape([1]*self.ndim)[tuple(index0)].shape

        # determine the interpolation axes
        # and for those axes, determine the lower index (inf) and the weight
        # (x) between lower and upper index
        interpolate_axis = []   # indices of the interpolated axes
        inf_list = []       # index of the lower elements for the interpolated axes
        x_list = []         # weights, for the interpolated axes
        for i in xrange(N):
            k = keys[i]

            # floating-point indices should be interpolated
            interpolate = False
            if isinstance(k, np.ndarray) and (k.dtype in [np.dtype('float32'),
                                                          np.dtype('float64')]):
                interpolate = True
                inf = k.astype('int')
                inf[inf == self.data.shape[i]-1] -= 1
                x = k-inf
                if k.ndim > 0:
                    x = x.reshape(shp_res)
            elif isinstance(k, float):
                interpolate = True
                inf = int(k)
                if inf == self.data.shape[i]-1:
                    inf -= 1
                x = k-inf
            if interpolate:
                # current axis needs interpolation
                inf_list.append(inf)
                x_list.append(x)
                interpolate_axis.append(i)

        # loop over the 2^n bracketing elements
        # (cartesian product of [0, 1] over n dimensions)
        n = len(interpolate_axis)
        result = 0
        for b in xrange(2**n):

            # coefficient attributed to the current item
            # and adjust the indices
            # for the interpolated dimensions
            coef = 1
            for i in xrange(n):
                # bb is the ith bit in b (0 or 1)
                bb = ((1<<i)&b)>>i
                x = x_list[i]
                if bb:
                    coef *= x
                else:
                    coef *= 1-x

                keys[interpolate_axis[i]] = inf_list[i] + bb

            result += coef * self.data[tuple(keys)]

        return result


    def equal(self, other, strict=True):
        '''
        Checks equality between two LUTs:
            - same axes
            - same shape
            - same values (if strict)
        '''
        if not isinstance(other, LUT):
            return False
        for i, ax in enumerate(self.axes):
            if (ax is None) and (other.axes[i] is None):
                continue
            if (ax is None) or (other.axes[i] is None):
                return False
            if not np.allclose(ax, other.axes[i]):
                return False
        if not self.data.shape == other.data.shape:
            return False
        if strict:
            if not np.allclose(self.data, other.data):
                return False
            if self.attrs != other.attrs:
                return False

        return True


    def __binary_operation_lut__(self, other, fn):
        '''
        apply fn(self, other) where other is a LUT
        the result is determined by using common axes between self and other
        and using appropriate broadcasting
        '''
        # shapes for broadcasting self vs other
        # None adds an singleton dimension, slice(None) adds a full dimension
        shp1, shp2 = [], []

        # new axes
        axes = []

        # determine union of axes
        names = interleave_seq(self.names, other.names)
        for i, a in enumerate(names):
            if a in self.names:
                axes.append(self.axes[self.names.index(a)])
                shp1.append(slice(None))
                if a in other.names:
                    shp2.append(slice(None))
                else:
                    shp2.append(None)
            else:
                axes.append(other.axes[other.names.index(a)])
                shp1.append(None)
                shp2.append(slice(None))

        # include common attributes
        attrs = {}
        for k in self.attrs:
            # check that the attributes are equal
            if not (k in other.attrs):
                continue
            if isinstance(self.attrs[k], np.ndarray):
                if not isinstance(other.attrs[k], np.ndarray):
                    continue
                if not np.allclose(self.attrs[k], other.attrs[k]):
                    continue
            else:
                if self.attrs[k] != other.attrs[k]:
                    continue
            attrs.update({k: self.attrs[k]})

        if self.desc == other.desc:
            desc = self.desc
        else:
            desc = str(fn)

        return LUT(
                fn(self.data[tuple(shp1)],
                    other.data[tuple(shp2)]),
                axes=axes, names=names,
                attrs=attrs, desc=desc)


    def __binary_operation_scalar__(self, other, fn):
        return LUT(fn(self.data, other),
                axes=self.axes, names=self.names,
                attrs=self.attrs, desc=self.desc)

    def __binary_operation__(self, other, fn):
        if isinstance(other, LUT):
            return self.__binary_operation_lut__(other, fn)
        else:
            return self.__binary_operation_scalar__(other, fn)

    def __add__(self, other):
        return self.__binary_operation__(other, lambda x, y: x+y)

    def __radd__(self, other):
        return self.__binary_operation__(other, lambda x, y: x+y)

    def __sub__(self, other):
        return self.__binary_operation__(other, lambda x, y: x-y)

    def __rsub__(self, other):
        return self.__binary_operation__(other, lambda x, y: y-x)

    def __mul__(self, other):
        return self.__binary_operation__(other, lambda x, y: x*y)

    def __rmul__(self, other):
        return self.__binary_operation__(other, lambda x, y: x*y)

    def __div__(self, other):
        return self.__binary_operation__(other, lambda x, y: x/y)

    def __rdiv__(self, other):
        return self.__binary_operation__(other, lambda x, y: y/x)

    def __truediv__(self, other):
        return self.__binary_operation__(other, lambda x, y: x/y)

    def __rtruediv__(self, other):
        return self.__binary_operation__(other, lambda x, y: y/x)

    def __eq__(self, other):
        return self.equal(other)

    def __neq__(self, other):
        return not self.equal(other)


    def to_mlut(self):
        '''
        convert to a MLUT
        '''
        m = MLUT()

        # axes
        if self.axes is not None:
            for i in xrange(len(self.axes)):
                name = self.names[i]
                axis = self.axes[i]
                if (name is None) or (axis is None):
                    continue
                m.add_axis(name, axis)

        # datasets
        m.add_dataset(self.desc, self.data, axnames=self.names, attrs=self.attrs)

        # attributes
        m.set_attrs(self.attrs)

        return m

    def apply(self, fn, desc=None):
        '''
        returns a LUT whose content is obtained by applying function fn
        if desc is provided, use this description
        '''
        if (desc is None) and (self.desc is not None):
            desc = self.desc
        return LUT(fn(self.data),
                axes=self.axes, names=self.names,
                attrs=self.attrs, desc=desc)

    def reduce(self, fn, axis, grouping=None, as_lut=False, **kwargs):
        '''
        apply function fn to a given axis
        fn: function to apply
            should be applicable to a numpy.ndarray and support argument axis
            (example: numpy.sum)
        axis: name (str) or index of axis
        grouping: iterable of same size as axis
                  fn is applied by groups corresponding to identical values in
                  grouping
                      example: grouping = [0, 0, 0, 1, 1, 2]
                      results in fn(3 first elements), fn(2 next), then fn(last)
                    the axis of the reduced axis takes the values of grouping
                  default None (apply to all elements, remove axis)
        as_lut: if axis reduction results in a scalar, returns a dimensionless LUT
                default False (returns a scalar)

        '''
        if isinstance(axis, str):
            index = self.names.index(axis)
        else:
            index = axis

        if grouping is None:
            axes = list(self.axes)
            names = list(self.names)
            axes.pop(index)
            names.pop(index)
            if (self.ndim == 1) and (not as_lut):
                # returns a scalar
                return fn(self.data, axis=index, **kwargs)
            else:
                # returns a LUT
                return LUT(fn(self.data, axis=index, **kwargs),
                        axes=axes, names=names,
                        attrs=self.attrs, desc=self.desc)
        else:
            assert len(grouping) == len(self.axes[index])
            shp = list(self.data.shape)
            U = uniq(grouping)
            shp[index] = len(U)
            data = np.zeros(shp, dtype=self.data.dtype)
            ind1 = [slice(None),] * self.ndim
            ind2 = [slice(None),] * self.ndim
            for i, u in enumerate(U):
                # fill each group
                ind1[index] = i
                ind2[index] = (grouping == u)
                data[tuple(ind1)] = fn(self.data[tuple(ind2)], axis=index, **kwargs)
            axes = list(self.axes)
            axes[index] = U
            return LUT(data,
                    axes=axes, names=self.names,
                    attrs=self.attrs, desc=self.desc)


    def swapaxes(self, axis1, axis2):
        '''
        Swaps two axes of the LUT

        Arguments:
            axis1 and axis2 (int or str): the name or index of the axes to swap

        Returns: the swapped LUT
        '''
        if isinstance(axis1, str):
            axis1 = self.names.index(axis1)
        if isinstance(axis2, str):
            axis2 = self.names.index(axis2)

        # swap the names and axes
        names, axes = [], []
        for i in xrange(self.ndim):
            names.append(self.names[i])
            axes.append(self.axes[i])
        names[axis1], names[axis2] = names[axis2], names[axis1]
        axes[axis1], axes[axis2] = axes[axis2], axes[axis1]

        return LUT(self.data.swapaxes(axis1, axis2),
                   axes=axes, names=names,
                   attrs=self.attrs, desc=self.desc)


    def plot(self, *args, **kwargs):
        '''
        Plot a 1 or 2-dimension LUT
        Note:
        in 2-dim, creates a new figure

        Arguments:
            * swap: swap x and y axes
            * show_grid
            * fmt: format for plotting 1d data

            Only in 2-dim:
            * vmin, vmax: color scale extent (default: from min/max)
            * cmap: colormap instance
            * index: index (or Idx instance) for plotting transect
        '''
        if self.ndim == 0:
            if self.desc is None:
                print(self.formatter.format(self.data))
            else:
                print(('{}: '+self.formatter).format(self.desc, self.data))
        elif self.ndim == 1:
            self.__plot_1d(*args, **kwargs)
        elif self.ndim == 2:
            self.__plot_2d(*args, **kwargs)
        else:
            self.plot_nd(*args, **kwargs)

        return self


    def __plot_1d(self, show_grid=True, swap=False, fmt=None, label=None, 
                  vmin=None, vmax=None, legend=False, plot=None, **kwargs):
        '''
        plot a 1-dimension LUT, returns self

        arguments:
            plot: custom plot function (defaults to plot)
                  example: plot=semilogy
        '''
        from pylab import xlabel, ylabel, grid, ylim, ticklabel_format
        import pylab as pl

        if plot is None:
            plot = pl.plot

        # no plotting for string datasets
        if self.data.dtype.char == 'S':
            warnings.warn('1D plot does not work for string arrays')
            return

        ax = self.axes[0]
        if ax is None:
            ax = range(self.shape[0])

        if vmin is None:
            vmin = np.amin(self.data[~np.isnan(self.data)])
        if vmax is None:
            vmax = np.amax(self.data[~np.isnan(self.data)])

        if not swap:
            xx = ax
            yy = self.data
            xlab = self.names[0]
            ylab = self.desc
        else:
            xx = self.data
            yy = ax
            xlab = self.desc
            ylab = self.names[0]

        if fmt is None:
            plot(xx, yy, label=label)
        else:
            plot(xx, yy, fmt, label=label)
        ylim(vmin,vmax)
        if xlab is not None:
            xlabel(xlab)
        if ylab is not None and not legend:
            ylabel(ylab)

        if plot == pl.plot:
            ticklabel_format(axis='y', style='sci', scilimits=(-2,2))

        grid(show_grid)
        if legend:
            leg = pl.legend(loc='best', fancybox=True)
            leg.get_frame().set_alpha(0.5)


    def __plot_2d(self, fmt='k-', show_grid=True, swap=False,
                  vmin=None, vmax=None, cmap=None, index=None,
                  label=None, **kwargs):
        '''
        Plot a 2-dimension LUT, with optional transect
        returns self
        '''
        import matplotlib.pyplot as plt

        if vmin is None:
            vmin = np.amin(self.data[~np.isnan(self.data)])
        if vmax is None:
            vmax = np.amax(self.data[~np.isnan(self.data)])

        if cmap is None:
            cmap = plt.cm.jet
            cmap.set_under('black')
            cmap.set_over('white')
            cmap.set_bad('0.5') # grey 50%

        if not swap:
            axis1, axis2 = self.axes[1], self.axes[0]
            lab1, lab2 = self.names[1], self.names[0]
            data = self.data
        else:
            axis1, axis2 = self.axes[0], self.axes[1]
            data = np.swapaxes(self.data, 0, 1)
            lab1, lab2 = self.names[0], self.names[1]

        if axis1 is None:
            axis1 = np.arange(data.shape[1])
        if axis2 is None:
            axis2 = np.arange(data.shape[0])

        if index is None:
            fig, ax1 = plt.subplots(nrows=1, ncols=1)
        else:
            fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, sharex=True)

        axis1 = bin_edges(axis1)
        axis2 = bin_edges(axis2)

        X, Y = np.meshgrid(axis1, axis2)
        plt.sca(ax1)
        im = plt.pcolormesh(X, Y, data, vmin=vmin, vmax=vmax)
        plt.grid(show_grid)
        plt.axis([
            np.amin(axis1),
            np.amax(axis1),
            np.amin(axis2),
            np.amax(axis2),
            ])
        cbar_ax = fig.add_axes([0.92, 0.25, 0.03, 0.5])
        # cbar_ax.text(0, -0.1, self.desc, verticalalignment='top')
        if label is None:
            label = self.desc
        if label is not None:
            cbar_ax.set_title(self.desc, weight='bold', horizontalalignment='left', position=(0.,-0.15))
        plt.colorbar(im, extend='both', orientation='vertical', cax=cbar_ax)

        if index is None:
            plt.sca(ax1)
            plt.ylabel(lab2)
            plt.xlabel(lab1)
        else:
            # show transect
            if isinstance(index, Idx):
                v = index.value
            else:
                v = axis2[index]

            ax1.plot([np.amin(axis1), np.amax(axis1)], [v, v], 'w--')
            ax1.ticklabel_format(axis='y', style='sci', scilimits=(-2,2))

            if not swap:
                ax2.plot(axis1, self[index, :], fmt)
            else:
                ax2.plot(axis1, self[:, index], fmt)
            ax2.grid(show_grid)

            plt.sca(ax1)
            plt.ylabel(lab2)
            plt.sca(ax2)
            plt.ylabel(self.desc)
            plt.xlabel(lab1)

            ax2.axis(ymin=vmin, ymax=vmax)
            ax2.ticklabel_format(axis='y', style='sci', scilimits=(-2,2))

        fig.subplots_adjust(hspace=0.)

    def plot_nd(self, *args, **kwargs):
        try:
            from ipywidgets import VBox, HBox, Checkbox, IntSlider, HTML
            from IPython.display import display, clear_output
        except ImportError:
            raise Exception('IPython notebook widgets are required '
                    'to plot a LUT with more than 2 dimensions')
        wid = []
        chks, sliders, texts = [], [], []
        for i in xrange(self.ndim):
            name = self.names[i]
            ax = self.axes[i]
            desc = name
            if desc is None:
                desc = 'NoName'
            if ax is not None:
                desc += ' [{:.5g}, {:.5g}]'.format(ax[0], ax[-1])
            chk = Checkbox(description=desc, value=False)
            slider = IntSlider(min=0, max=self.shape[i]-1)
            slider.visible = False
            text = HTML(value='')
            text.visible = False
            chks.append(chk)
            sliders.append(slider)
            texts.append(text)
            wid.append(HBox([chk, text, slider]))

        def update():
            clear_output()

            keys = []
            ndim = 0
            for i in xrange(self.ndim):
                # set sliders visibility
                sliders[i].visible = chks[i].value
                texts[i].visible = chks[i].value

                if chks[i].value:
                    index = sliders[i].value
                    keys.append(index)
                    # set text (value)
                    if self.axes[i] != None:
                        texts[i].value = '&nbsp;&nbsp;{:.5g}&nbsp;&nbsp;'.format(self.axes[i][index])
                else:
                    keys.append(slice(None))
                    ndim += 1

            if ndim <= 2:
                self.sub().__getitem__(tuple(keys)).plot(*args, **kwargs)
            else:
                print('Please select at least {} dimension(s)'.format(ndim-2))

        for i in xrange(self.ndim):
            chks[i].on_trait_change(update, 'value')
            sliders[i].on_trait_change(update, 'value')

        display(VBox(wid))
        update()

    def plot_polar(self, *args, **kwargs):
        plot_polar(self, *args, **kwargs)
        return self

    def plot_semi(self, *args, **kwargs):
        kwargs['semi'] = True
        plot_polar(self, *args, **kwargs)
        return self

    def transect2D(self, *args, **kwargs):
        transect2D(self, *args, **kwargs)
        return self
    
    def rename_axis(self, ax1, ax2):
        if self.names is None:
            raise Exception("Requested rename_axis on a LUT that has no named axes")
        else:
            self.names = [ax2 if x == ax1 else x for x in self.names]

        return self

    def to_xarray(self, deduplicate={}):
        """
        Convert LUT to `xr.DataArray`

        Arguments:
        ----------

        deduplicate: dictionary used to rename duplicate dimensions within the LUT
            ex: {'a': ['a0', 'a1']}
                rename duplicate dimension 'a' to 'a0', 'a1'
        """
        idim = itertools.count()
        new_dims = [x
                    if x is not None
                    else f'dim_{next(idim)}'
                    for x in self.names]
        for d in deduplicate:
            idedup = itertools.count()
            new_dims = [x
                        if x not in deduplicate
                        else deduplicate[x][next(idedup)]
                        for x in new_dims]

        da = xr.DataArray(
            self.data,
            dims=new_dims,
            )
        for i, name in enumerate(self.names):
            if (name is None) or (self.axes[i] is None):
                continue
            if name in deduplicate:
                for n in deduplicate[name]:
                    da = da.assign_coords(**{n: self.axes[i]})
            else:
                da = da.assign_coords(**{name: self.axes[i]})
        da.attrs = self.attrs
        return da


def Idx(value, name=None, round=False, fill_value=None):
    '''
    Calculate the indices of values by interpolation in a LUT axis
    The index method is typically called when indexing a dimension of a LUT
    object by a Idx object.
    The round attribute (boolean) indicates whether the resulting index should
    be rounded to the closest integer.

    Idx value can be of several kinds:
    - scalar (int or float)
     => index return scalar float index
    - array-like
      => index return array-like float index
    - callable (ex: lambda x: x<5)
      => index returns callable applied to axis

    Example: find the float index of 35. in an array [0, 10, ..., 100]
    >>> Idx(35.).index(np.linspace(0, 100, 11))
    array(3.5)

    Find the indices of several values in the array [0, 10, ..., 100]
    >>> Idx(np.array([32., 45., 72.])).index(np.linspace(0, 100, 11))
    array([ 3.2,  4.5,  7.2])

    Optionally, the name of the parameter can be provided as a keyword
    argument.
    Example: Idx(3., 'a') instead of Idx(3.)
    This allows verifying that the parameter is used in the right axis.

    Options:
        - fill_value are passed to interp1d (implies bounds_error=False)
          Special value for fill_value='extrema': fill with extrema values, don't extrapolate

    NOTE: this function is a factory that returns an Idx_* instance based on input type

    Method apply(axis):
        returns the Idx values applied to axis instead of their index
    '''
    if hasattr(value, '__call__'):
        return Idx_filter(value, name=name)
    else:
        return Idx_arr(value, name=name,
                       round=round, fill_value=fill_value)


class Idx_base(object):
    '''
    Base class for Idx objects
    '''
    def __init__(self, value, name=None, round=False, fill_value=None):
        if value is not None:
            self.value = value
            self.name = name
            self.round = round
            self.fill_value = fill_value


class Idx_arr(Idx_base):
    '''
    Idx class for array-like values
    '''
    def index(self, axis):
        '''
        Return the floating point index of the values in the axis
        '''
        if len(axis) == 1:
            if not np.allclose(np.array(self.value), axis[0]):
                raise ValueError("(Idx) Out of axis value (value={}, axis={})".format(self.value, axis))
            return 0

        else:
            # axis is scalar or ndarray: interpolate
            if self.fill_value == 'extrema':
                fv = (0, len(axis)-1)
            elif self.fill_value == 'extrema,warn':
                fv = (0, len(axis)-1)
                if (np.amax(self.value) > np.amax(axis)):
                    warnings.warn('(Idx) Value {} is above the axis maximum {} (axis {})'.format(
                        np.amax(self.value), np.amax(axis), self.name))
                if (np.amin(self.value) < np.amin(axis)):
                    warnings.warn('(Idx) Value {} is under the axis minimum {} (axis {})'.format(
                        np.amin(self.value), np.amin(axis), self.name))
            else:
                fv = self.fill_value
            be = (self.fill_value is None)

            res = interp1d(axis, np.arange(len(axis)),
                    bounds_error=be,
                    fill_value=fv)(self.value)
            if self.round:
                if isinstance(res, np.ndarray):
                    res = res.round().astype(int)
                else:
                    res = round(res)
            return res


    def apply(self, axis=None):
        return self.value

class Idx_filter(Idx_base):
    '''
    Idx class for filtering functions
    '''
    def index(self, axis):
        # value is a callable (function):
        # returns value applied to axis
        return self.value(axis)

    def apply(self, axis):
        return axis[self.index(axis)]


class Subsetter(object):
    '''
    A conveniency class to use the syntax like:
    LUT.sub()[:,:,0]
    for subsetting LUTs
    '''
    def __init__(self, LUT):
        self.LUT = LUT

    def __getitem__(self, keys):
        '''
        subset parent LUT
        '''
        return self.LUT.sub(dict(enumerate(keys)))


def plot_polar(lut, index=None, vmin=None, vmax=None, rect=211, sub=212,
               sym=True, swap='auto', fig=None, cmap=None, semi=False):
    '''
    Contour and eventually transect of 2D LUT on a semi polar plot, with
    dimensions (angle, radius)

    lut: 2D look-up table to display
            with axes (radius, angle) (unless swapped)
            angle is assumed to be in degrees and is not scaled
    index: index of the item to transect in the 'angle' dimension
           can be an Idx instance of several values or a list of indices
           if None (default), no transect
    vmin, vmax: range of values
                default None: determine min/max from values
    rect: subplot position of the main plot (111 for example)
    sub: subplot position of the transect
    sym: the transect uses symmetrical axis (boolean)
         if None (default), use symmetry iff axis is 'zenith'
    swap: swap the order of the 2 axes to (radius, angle)
          if 'auto', searches for 'azi' in both axes names
    fig : destination figure. If None (default), create a new figure.
    cmap: color map
    semi: polar by default, otherwise semi polar if lut is computed for 360 deg
    '''
    from pylab import figure, cm
    import mpl_toolkits.axisartist.angle_helper as angle_helper
    from matplotlib.transforms import Affine2D
    from mpl_toolkits.axisartist import floating_axes
    from matplotlib.projections import PolarAxes

    #
    # initialization
    #
    Phimax = 360.
    if semi : Phimax=180.

    assert lut.ndim == 2

    show_sub = index is not None
    if fig is None:
        if show_sub:
            fig = figure(figsize=(4.5, 4.5))
        else:
            fig = figure(figsize=(4.5, 6))

    if swap =='auto':
        if ('azi' in lut.names[1].lower()) and ('azi' not in lut.names[0].lower()):
            swap = True
        else:
            swap = False

    # ax1 is angle, ax2 is radius
    if swap:
        ax1, ax2 = lut.axes[1], lut.axes[0]
        name1, name2 = lut.names[1], lut.names[0]
        data = np.swapaxes(lut.data, 0, 1)
    else:
        ax1, ax2 = lut.axes[0], lut.axes[1]
        name1, name2 = lut.names[0], lut.names[1]
        data = lut.data

    if vmin is None:
        vmin = np.amin(lut.data[~np.isnan(lut.data)])
    if vmax is None:
        vmax = np.amax(lut.data[~np.isnan(lut.data)])
    if vmin == vmax:
        vmin -= 0.001
        vmax += 0.001
    if vmin > vmax: vmin, vmax = vmax, vmin

    #
    # semi polar axis
    #
    # angle
    ax1_scaled = ax1
    label1 = name1

    # radius axis
    # scale to [0, 90]
    ax2_min = np.amin(ax2)
    ax2_max = np.amax(ax2)
    ax2_scaled = (ax2 - ax2_min)/(ax2_max- ax2_min)*90.
    label2 = name2

    # angle axis
    grid_locator1 = angle_helper.LocatorDMS({True: 4, False:8}[semi], include_last=False)
    tick_formatter1 = angle_helper.FormatterDMS()

    class Locator(object):
        def __call__(self, *args):
            # as returned by angle_helper.step
            return [np.array([0, 30, 60, 90]), 4, 1.0]
    class Formatter(object):
        def __call__(self, *args):
            return map(lambda x: '{:.3g}'.format(x), np.linspace(ax2_min, ax2_max, 4))

    # radius axis
    if ((ax2_min < 10.) and (ax2_min >= 0)
            and (ax2_max <= 90) and (ax2_max > 80)):
        # angle in degrees
        grid_locator2 = angle_helper.LocatorDMS(4)
        tick_formatter2 = angle_helper.FormatterDMS()
    else:
        # custom locator
        grid_locator2 = Locator()
        tick_formatter2 = Formatter()

    tr_rotate = Affine2D().translate(0, 0)  # orientation
    tr_scale = Affine2D().scale(np.pi/180., 1.)  # scale to radians

    tr = tr_rotate + tr_scale + PolarAxes.PolarTransform()

    grid_helper = floating_axes.GridHelperCurveLinear(tr,
                                    extremes=(0., Phimax, 0., 90.),
                                    grid_locator1=grid_locator1,
                                    grid_locator2=grid_locator2,
                                    tick_formatter1=tick_formatter1,
                                    tick_formatter2=tick_formatter2,
                            )

    ax_polar = floating_axes.FloatingSubplot(fig, rect, grid_helper=grid_helper)
    fig.add_subplot(ax_polar)

    # adjust axis
    ax_polar.grid(True)
    ax_polar.axis["left"].set_axis_direction("bottom")
    ax_polar.axis["right"].set_axis_direction("top")
    ax_polar.axis["bottom"].set_visible(False)
    ax_polar.axis["top"].set_axis_direction("bottom")
    ax_polar.axis["top"].toggle(ticklabels=True, label=True)
    ax_polar.axis["top"].major_ticklabels.set_axis_direction("top")
    ax_polar.axis["top"].label.set_axis_direction("top")

    #ax_polar.axis["top"].axes.text(0.70, 0.92, label1,
    ax_polar.axis["top"].axes.text(0.72, 0.98, label1,
                                   transform=ax_polar.transAxes,
                                   ha='left',
                                   va='bottom')
    #ax_polar.axis["left"].axes.text(0.25, -0.03, label2,
    ax_polar.axis["left"].axes.text(0.10, -0.03, label2,
                                   transform=ax_polar.transAxes,
                                   ha='center',
                                   va='top')

    # create a parasite axes whose transData in RA, cz
    aux_ax_polar = ax_polar.get_aux_axes(tr)

    aux_ax_polar.patch = ax_polar.patch # for aux_ax to have a clip path as in ax
    ax_polar.patch.zorder=0.9 # but this has a side effect that the patch is
                        # drawn twice, and possibly over some other
                        # artists. So, we decrease the zorder a bit to
                        # prevent this.

    #
    # initialize the cartesian axis below the semipolar
    #
    if show_sub:
        ax_cart = fig.add_subplot(sub)
        if sym:
            ax_cart.set_xlim(-ax2_max, ax2_max)
        else:
            ax_cart.set_xlim(ax2_min, ax2_max)
        ax_cart.set_ylim(vmin, vmax)
        ax_cart.ticklabel_format(axis='y', style='sci', scilimits=(-2,2))
        ax_cart.grid(True)

    #
    # draw colormesh
    #
    if cmap is None:
        cmap = cm.rainbow.copy()
        cmap.set_under('black')
        cmap.set_over('white')
        cmap.set_bad('0.5') # grey 50%
    r, t = np.meshgrid(bin_edges(ax2_scaled, min=0, max=90), bin_edges(ax1_scaled))
    masked_data = np.ma.masked_where(np.isnan(data) | np.isinf(data), data)
    im = aux_ax_polar.pcolormesh(t, r, masked_data, cmap=cmap, vmin=vmin, vmax=vmax)

    if show_sub:
        # convert Idx instance to index if necessarry
        if isinstance(index, Idx_base):
            indexes = np.round(index.index(ax1)).astype(int)
            #indexes = int(round(index.index(ax1)))
        else: indexes = index
        for ii,index in enumerate(indexes):
            if semi:
                mirror_index = -1 -index
            else:
                mirror_index = (ax1_scaled.shape[0]//2 + index)%ax1_scaled.shape[0]
            # draw line over colormesh
            vertex0 = np.array([[0,0],[ax1_scaled[index],ax2_max]])
            vertex1 = np.array([[0,0],[ax1_scaled[mirror_index],ax2_max]])
            aux_ax_polar.plot(vertex0[:,0],vertex0[:,1], 'w')
            if sym:
                aux_ax_polar.plot(vertex1[:,0],vertex1[:,1],'w--',linewidth=2)

            #
            # plot transects
            #
            color = ['k', 'r','g','b','m','y'][ii%6]
            ax_cart.plot(ax2, data[index,:],'-'+color)
            if sym:
                ax_cart.plot(-ax2, data[mirror_index,:],'--'+color)

    # add colorbar
    fig.colorbar(im, orientation='horizontal',
                 extend='both', ticks=np.linspace(vmin, vmax, 5),
                 shrink=0.7)
    if lut.desc is not None:
        ax_polar.set_title(lut.desc, weight='bold', position=(0.05,0.97))


def transect2D(lut, index=None, vmin=None, vmax=None, sym=True, swap='auto', fig=None, sub=121, color='k', percent=False, fmt='-'):
    '''
    Transect of 2D LUT

    lut: 2D look-up table to display
            with axes (radius, angle) (unless swapped)
            angle is assumed to be in degrees and is not scaled
    index: index of the item to transect in the 'angle' dimension
           can be an Idx instance
           if None (default), no transect
    vmin, vmax: range of values
                default None: determine min/max from values
    sym: the transect uses symmetrical axis (boolean)
         if None (default), use symmetry if axis is 'zenith'
    swap: swap the order of the 2 axes to (radius, angle)
          if 'auto', searches for 'azi' in both axes names
    fig : destination figure. If None (default), create a new figure.
    color : color of the transect
    percent: if True set scale to 0 to 100%
    '''
    from pylab import figure

    assert lut.ndim == 2

    if fig is None:
        fig = figure(figsize=(4.5, 2.5))

    if swap =='auto':
        if ('azi' in lut.names[1].lower()) and ('azi' not in lut.names[0].lower()):
            swap = True
        else:
            swap = False

    # ax1 is angle, ax2 is radius
    if swap:
        ax1, ax2 = lut.axes[1], lut.axes[0]
        name1, name2 = lut.names[1], lut.names[0]
        data = np.swapaxes(lut.data, 0, 1)
    else:
        ax1, ax2 = lut.axes[0], lut.axes[1]
        name1, name2 = lut.names[0], lut.names[1]
        data = lut.data


    if vmin is None:
        vmin = np.amin(lut.data[~np.isnan(lut.data)])
    if vmax is None:
        vmax = np.amax(lut.data[~np.isnan(lut.data)])
    if vmin == vmax:
        vmin -= 0.001
        vmax += 0.001
    if vmin > vmax: vmin, vmax = vmax, vmin
    if percent:
        vmin=0.
        vmax=100.

    #
    # semi polar axis
    #
    # angle
    ax1_scaled = ax1
    label2 = name2

    # convert Idx instance to index if necessarry
    if isinstance(index, Idx_base):
        index = int(np.around(index.index(ax1)))
    mirror_index = (ax1_scaled.shape[0]//2 + index)%ax1_scaled.shape[0]

    if swap:
        title=lut.axes[1][index]
    else:
        title=lut.axes[0][index]

    # radius axis
    # scale to [0, 90]
    ax2_min = np.amin(ax2)
    ax2_max = np.amax(ax2)
    label1 = name1 + ' {:7.2f}'.format(title)
    #

    ax_cart = fig.add_subplot(sub)
    ax_cart.grid(True)

    if sym:
        ax_cart.set_xlim(-ax2_max, ax2_max)
    else:
        ax_cart.set_xlim(ax2_min, ax2_max)
    ax_cart.set_ylim(vmin, vmax)
    ax_cart.ticklabel_format(axis='y', style='sci', scilimits=(-2,2))
    ax_cart.grid(True)
    ax_cart.set_xlabel(label2)
    #ax_cart.set_title(label1)

    #
    # plot transects
    #
    ax_cart.plot(ax2, data[index,:],fmt, color=color)
    if sym:
       ax_cart.plot(-ax2, data[mirror_index,:],fmt, color=color)

    if lut.desc is not None:
        ax_cart.set_title(lut.desc)


def merge(M, axes, dtype=None):
    '''
    Merge several luts

    Arguments:
        - M is a list of MLUT objects to merge
        - axes is a list of axes names to merge
          these names should be present in each LUT attribute
        - dtype is the data type of the new axes
          ex: dtype=float
          if None, no data type conversion

    Returns a MLUT for which each dataset has new axes as defined in list axes
    (list of strings)
    The attributes of the merged mlut consists of the common attributes with
    identical values.

    >>> np.random.seed(0)

    Example: merge two MLUTS
    (also using attribute to dataset promotion)
    >>> M = []
    >>> for b in range(4):
    ...     M1 = MLUT()
    ...     M1.add_axis('ax1', np.arange(4))
    ...     M1.add_axis('ax2', np.arange(5)+10)
    ...     M1.add_dataset('a', np.random.randn(4, 5), ['ax1', 'ax2'])
    ...     M1.set_attrs({'b':b, 'c':b*10})
    ...     M1.promote_attr('b')  # attribute 'b' is converted to a scalar dataset
    ...     M.append(M1)
    >>> merged = merge(M, ['c'])
    >>> _=merged.print_info(show_self=False)
     Datasets:
      [0] a (float64 between -2.55 and 2.27), axes=('c', 'ax1', 'ax2')
      [1] b (float64 between 0 and 3), axes=('c',)
     Axes:
      [0] ax1: 4 values between 0 and 3
      [1] ax2: 5 values between 10 and 14
      [2] c: 4 values between 0 and 30

    '''

    m = MLUT()
    first = M[0]

    # check mluts compatibility
    for i in xrange(1, len(M)):
        assert first.equal(M[i], content=False, attributes=False, show_diff=True)

    # add old axes
    for (axname, axis) in first.axes.items():
        m.add_axis(axname, axis)

    # determine the new axes from the attributes of all mluts
    newaxes = []  # new axes
    newaxnames = []
    for axname in axes:
        axis = []
        for mlut in M:
            value = mlut.attrs[axname]
            if dtype is not None:
                value = dtype(value)
            if value not in axis:
                axis.append(value)
        m.add_axis(axname, axis)
        newaxes.append(axis)
        newaxnames.append(axname)

    # dataset loop
    for name in first.datasets():

        # build new data
        new_shape = tuple(map(len, newaxes))+first[name].shape
        _dtype = first[name].data.dtype
        newdata = np.zeros(new_shape, dtype=_dtype)
        try:
            newdata += np.nan
        except:
            pass
        for mlut in M:

            # find the index of the attributes in the new LUT
            index = ()
            for j, a in enumerate(axes):
                value = mlut.attrs[a]
                if dtype is not None:
                    value = dtype(value)
                index += (newaxes[j].index(value),)

            if first[name].data.ndim != 0:
                index += (slice(None),)

            newdata[index] = mlut[name].data

        axnames = first[name].names
        if axnames is None:
            m.add_dataset(name, newdata)
        else:
            m.add_dataset(name, newdata, newaxnames+axnames)

    # fill with common arguments
    for k, v in first.attrs.items():
        if False in map(lambda x: k in x.attrs, M):
            continue
        if isinstance(v, np.ndarray):
            if False in map(lambda x: np.allclose(v, x.attrs[k]), M):
                continue
        else:
            if False in map(lambda x: v == x.attrs[k], M):
                continue
        m.set_attr(k, v)

    return m


class MLUT(object):
    '''
    A class to store and manage multiple look-up tables

    How to create a MLUT:
    >>> m = MLUT()
    >>> m.add_axis('a', np.linspace(100, 150, 5))
    >>> m.add_axis('b', np.linspace(5, 8, 6))
    >>> m.add_axis('c', np.linspace(0, 1, 7))
    >>> np.random.seed(0)
    >>> m.add_dataset('data1', np.random.randn(5, 6), ['a', 'b'])
    >>> m.add_dataset('data2', np.random.randn(5, 6, 7), ['a', 'b', 'c'])
    >>> # Add a dataset without associated axes
    >>> m.add_dataset('data3', np.random.randn(10, 12))
    >>> m.set_attr('x', 12)   # set MLUT attributes
    >>> m.set_attrs({'y':15, 'z':8})
    >>> _=m.print_info(show_self=False)
     Datasets:
      [0] data1 (float64 between -2.55 and 2.27), axes=('a', 'b')
      [1] data2 (float64 between -2.22 and 2.38), axes=('a', 'b', 'c')
      [2] data3 (float64 between -2.77 and 2.3), axes=(None, None)
     Axes:
      [0] a: 5 values between 100.0 and 150.0
      [1] b: 6 values between 5.0 and 8.0
      [2] c: 7 values between 0.0 and 1.0

    Use bracket notation to extract a LUT
    Note that you can use a string or integer.
    data1 is the first dataset in this case, we could use m[0]
    >>> _=m['data1'].print_info()  # or m[0]
    LUT "data1" (float64 between -2.55 and 2.27):
      Dim 0 (a): 5 values betweeen 100.0 and 150.0
      Dim 1 (b): 6 values betweeen 5.0 and 8.0
    '''
    def __init__(self):
        # axes
        self.axes = OrderedDict()
        # data: a list of (name, array, axnames, attributes)
        self.data = []
        # attributes
        self.attrs = OrderedDict()

    def datasets(self):
        ''' returns a list of the datasets names '''
        return [x[0] for x in self.data]

    def add_axis(self, name, axis):
        ''' Add an axis to the MLUT '''
        assert isinstance(name, str)
        assert name not in self.axes, 'Axis "{}" already in MLUT'.format(name)
        if isinstance(axis, list):
            ax = np.array(axis)
        else:
            ax = axis
        assert ax.ndim == 1

        self.axes[name] = ax

    def add_dataset(self, name, dataset, axnames=None, attrs={}):
        '''
        Add a dataset to the MLUT
        name (str): name of the dataset
        dataset (np.array)
        axnames: list of (strings or None), or None
        attrs: dataset attributes
        '''
        assert name not in [x[0] for x in self.data], 'Error, "{}" already in MLUT'.format(name)
        if axnames is not None:
            # check axes consistency
            assert len(axnames) == len(dataset.shape)
            for i, ax in enumerate(axnames):
                if ax is None: continue
                if ax not in self.axes: continue
                assert dataset.shape[i] == len(self.axes[ax])
        else:
            axnames = [None]*dataset.ndim

        self.data.append((name, dataset, axnames, attrs))

    def add_lut(self, lut, desc=None):
        '''
        Add a LUT to the MLUT

        returns self
        '''
        if desc is None:
            desc = lut.desc
        assert desc is not None
        for iax in xrange(lut.ndim):
            axname = lut.names[iax]
            ax = lut.axes[iax]
            if axname in self.axes:
                # check axis
                if ax is not None:
                    assert np.array(self.axes[axname]).shape == np.array(ax).shape, \
                            'Inconsistent shapes for axis "{}": {} != {}'.format(
                                    axname, self.axes[axname].shape, ax.shape)
                    assert np.allclose(self.axes[axname], ax)
            elif axname is None:
                assert ax is None
            elif ax is not None:
                # add the axis
                self.add_axis(axname, ax)

        self.add_dataset(desc, lut.data, axnames=lut.names, attrs=lut.attrs)

        return self

    def rm_lut(self, name):
        ''' remove a LUT '''
        assert isinstance(name, (str, bytes))

        try:
            index = [x[0] for x in self.data].index(name)
        except ValueError:
            raise Exception('{} is not in {}', name, self)

        self.data.pop(index)

    def sub(self, d):
        '''
        The MLUT equivalent of LUT.sub

        returns a MLUT where each LUT is subsetted using dictionary d
        keys should only be strings
        values can be int, float, slice, Idx, array (bool or int)
        '''
        m = MLUT()
        for dd in d:
            assert isinstance(dd, str)

        for dd in self.datasets():
            m.add_lut(self[dd].sub(d, ignore=True))

        m.attrs = self.attrs

        return m

    def save(self, filename, fmt=None, overwrite=False,
             verbose=False, compress=True):
        '''
        Save a MLUT to filename
        fmt: output format: hdf4, netcdf4,
             or None (determine from filename extension)
        '''

        if exists(filename):
            if overwrite:
                remove(filename)
            else:
                ex = Exception('File {} exists'.format(filename))
                setattr(ex, 'filename', filename)
                raise ex

        if verbose:
            print('Writing "{}" to "{}" ({} format)'.format(self, filename, fmt))

        if fmt is None:
            if filename.endswith('.hdf'):
                fmt = 'hdf4'
            elif filename.endswith('.nc'):
                fmt = 'netcdf4'
            else:
                raise ValueError('Cannot determine desired format '
                        'of filename "{}"'.format(filename))

        if fmt=='netcdf4':
            self.__save_netcdf4(filename, overwrite=overwrite,
                              verbose=verbose, compress=compress)
        elif fmt=='hdf4':
            self.__save_hdf(filename, overwrite=overwrite,
                          verbose=verbose, compress=compress)
        else:
            raise ValueError('Invalid format {}'.format(fmt))

    def __save_netcdf4(self, filename, overwrite=False,
                      verbose=False, compress=True):
        from netCDF4 import Dataset
        root = Dataset(filename, 'w', format='NETCDF4')
        dummycount = 0

        # write axes and create associated dimensions
        for axname, ax in self.axes.items():
            root.createDimension(axname, len(ax))
            var = root.createVariable(axname, ax.dtype, [axname], zlib=compress)
            var[:] = ax[:]

        # write datasets
        for (name, data, axnames, attributes) in self.data:
            if verbose:
                print('   Write data "{}" ({}, {})'.format(name, data.dtype, data.shape))
            # create dummy dimensions when axis is missing
            for i in xrange(len(axnames)):
                if axnames[i] is None:
                    dummycount += 1
                    axnames[i] = 'dummy{:d}'.format(dummycount)
                if axnames[i] not in root.dimensions:
                    if verbose:
                        print('   Create dimension {}'.format(axnames[i]))
                    root.createDimension(axnames[i], data.shape[i])

            var = root.createVariable(name, data.dtype, axnames, zlib=compress)
            var.setncatts(attributes)
            var[:] = data[...]

        # write global attributes
        if verbose:
            print('   Write {} attributes'.format(len(self.attrs)))
        root.setncatts(self.attrs)

        root.close()


    def __save_hdf(self, filename, overwrite=False, verbose=False, compress=True):
        '''
        Save a MLUT to a hdf file
        '''
        from pyhdf.SD import SD, SDC
        from pyhdf.error import HDF4Error

        def safecast(data):
            # hdf4 does not support int64, uint64
            # cast to 32-bits
            if data.dtype == np.dtype('uint64'):
                assert np.allclose(data, data.astype('uint32'))
                return data.astype('uint32')
            if data.dtype == np.dtype('int64'):
                assert np.allclose(data, data.astype('int32'))
                return data.astype('int32')
            return data


        typeconv = {
                    np.dtype('float32'): SDC.FLOAT32,
                    np.dtype('float64'): SDC.FLOAT64,
                    np.dtype('uint64'): SDC.UINT32, # /!\
                    np.dtype('int64'): SDC.INT32,   # hdf4 does not support 64-bit ints
                    np.dtype('uint32'): SDC.UINT32,
                    np.dtype('int32'): SDC.INT32,
                    np.dtype('uint16'): SDC.UINT16,
                    np.dtype('int16'): SDC.INT16,
                    np.dtype('uint8'): SDC.UINT8,
                    np.dtype('int8'): SDC.INT8,
                    }
        hdf = SD(filename, SDC.WRITE | SDC.CREATE)

        # write axes
        if self.axes is not None:
            for name, ax in self.axes.items():
                if verbose:
                    print('   Write axis "{}" in "{}"'.format(name, filename))
                type = typeconv[ax.dtype]
                sds = hdf.create(name, type, ax.shape)
                if compress:
                    sds.setcompress(SDC.COMP_DEFLATE, 9)
                sds[:] = safecast(ax)[:]
                sds.endaccess()

        # write datasets
        for name, data, axnames, attrs in self.data:

            data = safecast(data)

            if verbose:
                print('   Write data "{}" ({}, {})'.format(name, data.dtype, data.shape))
            type = typeconv[data.dtype]

            # write data
            if data.ndim == 0:
                # scalar
                sds = hdf.create(name, type, (1,))
                setattr(sds, 'lut:scalar', 'True')
            else:
                sds = hdf.create(name, type, data.shape)
            if compress:
                sds.setcompress(SDC.COMP_DEFLATE, 9)
            sds[:] = data[...]

            if axnames not in [None, []]:
                setattr(sds, 'dimensions', ','.join(map(str, axnames)))
            if 'dimensions' in attrs:
                raise Exception('Error writing {}, "dimensions" attribute conflict'.format(filename))
            for k, v in attrs.items():
                setattr(sds, k, v)
            sds.endaccess()

        # write attributes
        if verbose:
            print('   Write {} attributes'.format(len(self.attrs)))
        for k, v in self.attrs.items():
            try:
                setattr(hdf, k, v)
            except HDF4Error:
                setattr(hdf, str(k), str(v))

        hdf.end()

    def set_attr(self, key, value):
        '''
        Set one attribute key -> value
        '''
        self.attrs[key] = value

        return self

    def set_attrs(self, attributes):
        '''
        Set multiple attributes to attrs
        attributes: dict
        '''
        self.attrs.update(attributes)

        return self

    def print_info(self, *args, **kwargs):
        # same as describe()
        return self.describe(*args, **kwargs)

    def describe(self, show_range=True, show_self=True, show_attrs=False, show_shape=False, show_axes=True, mem=False):
        total_size = 0
        if show_self:
            print(str(self))
        print(' Datasets:')
        for i, (name, dataset, axes, attrs) in enumerate(self.data):
            axdesc = ''
            if (axes is not None) and show_axes:
                axdesc += ', axes='+ str(tuple(axes))
            if show_shape:
                axdesc += ', shape={}'.format(dataset.shape)
            if show_range and isinstance(dataset, np.ndarray):
                try:
                    rng = ' in [{:.3g}, {:.3g}]'.format(np.amin(dataset), np.amax(dataset))
                except:
                    rng = ''
            else:
                rng = ''
            if mem:
                memdesc = ', {}'.format(sizeof_fmt(dataset.nbytes))
            else:
                memdesc = ''
            print('  [{}] {} ({}{})'.format(i, name, dataset.dtype, rng, dataset.shape) + axdesc + memdesc)
            total_size += dataset.nbytes

            if show_attrs and (len(attrs) != 0):
                print('    Attributes:')
                for k, v in attrs.items():
                    print('      {}: {}'.format(k, v))
        print(' Axes:')
        for i, (name, values) in enumerate(self.axes.items()):
            print('  [{}] {}: {} values in [{}, {}]'.format(i, name, len(values), values[0], values[-1]))
        if show_attrs:
            print(' Attributes:')
            for k, v in self.attrs.items():
                print(' ', k, ':', v)
        if mem:
            print('Total memory usage: {}'.format(sizeof_fmt(total_size)))

        return self

    def dropaxis(self, *ax_to_drop):
        '''
        drop axes of size 1

        example:
            m.dropaxis('a')   # remove axis 'a' of size 1
            m.dropaxis('a', 'b')   # remove 'a' and 'b'
        '''
        m = MLUT()

        # axes
        for axname, ax in self.axes.items():
            if axname not in ax_to_drop:
                m.add_axis(axname, ax)

        # datasets
        for (name, data, axnames, attributes) in self.data:
            axes = [x for x in axnames if x not in ax_to_drop]
            shp = [s for (x, s) in zip(axnames, data.shape) if x not in ax_to_drop]
            m.add_dataset(name, data.reshape(shp), axnames=axes, attrs=attributes)

        # global attributes
        m.attrs = self.attrs

        return m


    def promote_attr(self, name):
        '''
        Create a new dataset from attribute name
        '''
        assert isinstance(name, str)
        assert name in self.attrs
        value = np.array(self.attrs[name])

        self.add_dataset(name, value)

    def __getitem__(self, key):
        '''
        return the LUT corresponding to key (int or string)
        '''
        if isinstance(key, (str, unicode)):
            index = -1
            for i, (name, _, _, _) in enumerate(self.data):
                if key == name:
                    index = i
                    break
            if index == -1:
                raise Exception('Cannot find dataset {}'.format(key))
        elif isinstance(key, int):
            index = key
        else:
            raise Exception('multi-dimensional LUTs should only be indexed with strings or integers')

        name, dataset, axnames, attrs = self.data[index]
        if axnames is None:
            axes = None
        else:
            axes = []
            for ax in axnames:
                if (ax is None) or (ax not in self.axes):
                    axes.append(None)
                else:
                    axes.append(self.axes[ax])

        return LUT(desc=name, data=dataset, axes=axes, names=axnames, attrs=attrs)

    def equal(self, other, content=True, attributes=True, show_diff=False):
        '''
        Test equality between two MLUTs
        Arguments:
         * show_diff: print their differences
         * content: check LUT content (otherwise only axes and shapes)
         * attributes: check global attributes
        '''
        msg = 'MLUTs diff:'
        if not isinstance(other, MLUT):
            msg += '  other is not a MLUT ({})'.format(str(other))
            print(msg)
            return False

        eq = True

        # check axes
        for k in set(self.axes).union(other.axes):
            if (k not in other.axes) or (k not in self.axes):
                msg += '  axis {} missing in either\n'.format(k)
                eq = False
            if self.axes[k].shape != other.axes[k].shape:
                msg += '  axis {} shape mismatch\n'.format(k)
                eq = False
            if not np.allclose(self.axes[k], other.axes[k]):
                msg += '  axis {} is different\n'.format(k)
                eq = False

        # check datasets
        if set(self.datasets()) != set(other.datasets()):
            msg += '  Datasets are different\n'
            msg += '   -> {}'.format(str(self.datasets()))
            msg += '   -> {}'.format(str(other.datasets()))
            eq = False
        for name in self.datasets():
            if not self[name].equal(other[name], strict=content):
                msg += '  dataset {} differs (shapes are {} and {})\n'.format(name, self[name].shape, other[name].shape)
                eq = False

        # check global attributes
        if attributes:
            for a in set(self.attrs.keys()).union(other.attrs.keys()):
                if (a not in self.attrs) or (a not in other.attrs):
                    msg += '  attribute {} missing in either MLUT\n'.format(a)
                    eq = False
                    continue
                if (self.attrs[a] != other.attrs[a]):
                    msg += '  value of attribute {} differs ({} and {})\n'.format(a, self.attrs[a], other.attrs[a])
                    eq = False
                    continue

        if show_diff and not eq:
            print(msg)

        return eq

    def __eq__(self, other):
        return self.equal(other)

    def __neq__(self, other):
        return not self.equal(other)

    def axis(self, axname, aslut=False):
        '''
        returns an axis
        if aslut: returns it as a LUT
        otherwise, as values
        '''
        data = self.axes[axname]
        if aslut:
            return LUT(desc=axname, data=data, axes=[data], names=[axname])
        else:
            return data

    def plot(self, datasets=None, extra_widgets=True, *args, **kwargs):
        '''
        display all datasets in the MLUT
        * datasets: list of datasets to display
                    if None, display all datasets
        * extra_widgets: show extra widgets for:
                 1) interactive selection of datasets to display
                 2) choice of min/max values
        '''
        try:
            from ipywidgets import VBox, HBox, Checkbox, IntSlider, HTML, FloatText, Button
            from pylab import axis
            from IPython.display import display, clear_output
        except ImportError:
            raise Exception('IPython notebook widgets are required '
                    'to plot a MLUT')

        if datasets is None:
            datasets = self.datasets()

        wid = []
        axes = {}  # name: ()
        dstchk = {}  # dataset checkboxes
        if 'vmin' in kwargs:
            vmin = kwargs['vmin']
        else:
            vmin = 0
        if 'vmax' in kwargs:
            vmax = kwargs['vmax']
        else:
            vmax = 1
        vminmax = [   # vmin/vmax/xmin/xmax float texts
                FloatText(description='vmin', value=vmin),
                FloatText(description='vmax', value=vmax),
                Checkbox(description='xmin/xmax', value=False),
                FloatText(description='xmin', value=0),
                FloatText(description='xmax', value=1),
                ]
        for d in datasets:
            dstchk[d] = Checkbox(description=d, value=False)
            for iax, name in enumerate(self[d].names):
                if name is None:
                    raise Exception('Does not work with unnamed axes, sorry :/')
                if name in axes:
                    continue

                ax = self[d].axes[iax]
                desc = name
                if ax is not None:
                    desc += ' [{:.5g}, {:.5g}]'.format(ax[0], ax[-1])

                chk = Checkbox(description=desc, value=True)
                slider = IntSlider(min=0, max=self[d].shape[iax]-1)
                slider.visible = False
                text = HTML(value='')
                text.visible = False
                button_minus = Button(description='-')
                setattr(button_minus, 'slider', slider)
                button_plus  = Button(description='+')
                setattr(button_plus, 'slider', slider)

                wid.append(HBox([chk, text, slider, button_minus, button_plus]))

                axes[name] = (chk, slider, text, button_minus, button_plus)

        if extra_widgets:
            wid.insert(0, HTML(value='<b>AXES:</b>'))
            wid.insert(0, HBox(dstchk.values()))
            wid.insert(0, HTML(value='<b>DATASETS:</b>'))
            wid.insert(0, HBox(vminmax))

        def update():
            clear_output()

            # update vmin/vmax
            if extra_widgets:
                kwargs['vmin'] = vminmax[0].value
                kwargs['vmax'] = vminmax[1].value
                vminmax[3].visible = vminmax[2].value
                vminmax[4].visible = vminmax[2].value

            # update sliders visibility
            for a, (chk, slider, text, button_minus, button_plus) in axes.items():
                slider.visible = chk.value
                text.visible = chk.value
                button_minus.visible = chk.value
                button_plus.visible = chk.value

            # display each dataset
            for d in datasets:

                if not dstchk[d].value:
                    continue

                keys = []
                ndim = 0

                for i, name in enumerate(self[d].names):
                    chk, slider, text, _, _ = axes[name]

                    if chk.value:
                        index = slider.value
                        keys.append(index)
                        # set text (value)
                        if self[d].axes[i] != None:
                            text.value = '&nbsp;&nbsp;{:.5g}&nbsp;&nbsp;'.format(self[d].axes[i][index])
                    else:
                        keys.append(slice(None))
                        ndim += 1

                if ndim <= 2:
                    self[d].sub().__getitem__(tuple(keys)).plot(label=d, legend=True, *args, **kwargs)
                    if vminmax[2].value:
                        axis(xmin=vminmax[3].value, xmax=vminmax[4].value)
                else:
                    print('{}: Please select at least {} dimension(s)'.format(d, ndim-2))
        def decrement(b):
            b.slider.value -= 1
        def increment(b):
            b.slider.value += 1
        for a, (chk, slider, text, button_minus, button_plus) in axes.items():
            chk.on_trait_change(update, 'value')
            slider.on_trait_change(update, 'value')
            button_minus.on_click(decrement)
            button_plus.on_click(increment)
        for fl in vminmax:
            fl.on_trait_change(update, 'value')
        for chk in dstchk.values():
            chk.on_trait_change(update, 'value')

        display(VBox(wid))
        update()

    def rename_axis(self, ax1, ax2):
        '''
        Rename axis ax1 to ax2 (in place)
        '''
        # modify axes
        self.axes = OrderedDict(((ax2 if k == ax1 else k, v) for k, v in self.axes.items()))

        # modify data
        self.data = [(name,
                      array,
                      None if axnames is None else [ax2 if x == ax1 else x for x in axnames],
                      attributes)
                     for (name, array, axnames, attributes) in self.data]

        return self

    def to_xarray(self):
        """
        Convert LUT to `xr.Dataset`
        """
        ds = xr.Dataset()

        idim = itertools.count()

        for (name, array, axnames, attributes) in self.data:
            ds[name] = xr.DataArray(
                array,
                dims=[x if x is not None else f'dim_{next(idim)}' for x in axnames],
            )
            ds[name].attrs = attributes
        
        ds = ds.assign_coords(**self.axes)
        ds.attrs = self.attrs

        return ds


def read_mlut(filename, fmt=None):
    '''
    Read a MLUT (multi-format)
    fmt: netcdf4, hdf4
         or None (determine format from extension)
    '''
    if fmt is None:
        if filename.endswith('.hdf'):
            fmt = 'hdf4'
        elif filename.endswith('.nc'):
            fmt = 'netcdf4'
        else:
            raise ValueError('Cannot determine desired format '
                    'of filename "{}"'.format(filename))

    if fmt=='netcdf4':
        return read_mlut_netcdf4(filename)
    elif fmt=='hdf4':
        return read_mlut_hdf(filename)
    elif fmt=='hdf5':
        return read_mlut_hdf5(filename)

    else:
        raise ValueError('Invalid format {}'.format(fmt))


def read_mlut_netcdf4(filename):
    '''
    Read a MLUT (netcdf4 format)
    '''
    # assumes everything is in the root group
    m = MLUT()
    from netCDF4 import Dataset
    root = Dataset(filename, 'r', format='NETCDF4')

    # read axes
    for dim in root.dimensions:
        if (not dim.startswith('dummy')) and (dim in root.variables):
            m.add_axis(str(dim), filled(root.variables[dim][:]))

    # read datasets
    for varname in root.variables:
        if varname in m.axes:
            continue
        var = root.variables[varname]

        # read attributes
        attrs = {}
        for a in var.ncattrs():
            attrs[a] = var.getncattr(a)

        m.add_dataset(varname, filled(var[:]), [str(x) for x in var.dimensions], attrs=attrs)

    # read global attributes
    for a in root.ncattrs():
        m.set_attr(a, root.getncattr(a))

    root.close()

    return m


def read_mlut_hdf5(filename, datasets=None, lazy=False, group=None, wrap_data=None):
    '''
    read a MLUT from a hdf5 file (filename)
    
    - datasets: list of datasets to read:
        * None (default): read all datasets, including axes as indicated by the
          attribute 'dimensions'
        * a list of:
            - dataset names (string)
            - or a tuple (dataset_name, axes) where axes is a list of
              dimensions (strings), overriding the attribute 'dimensions'
    - lazy: whether to activate lazy reading (data is returned as hdf object)
    - wrap_data: a function applied to each data. Can be dask.array.array to
      return as dask arrays, which is useful for further conversion to xarray Dataset
    '''
    import h5py

    ff = h5py.File(filename,"r")

    if group:
        f = ff[group]
    else:
        f = ff

    # set the list of dataset
    if datasets is None:
        ls_datasets = list(f['data'].keys())
    else:
        ls_datasets = datasets

    # look for axis rquired for the datasets
    ls_axis   = []
    axis_data = []
    for dataset in ls_datasets: 
        if dataset in f['data'].keys():
            if not f['data'][dataset].attrs.__contains__('dimensions'):
                print('Missing -dimensions- Attr in dataset "{}" '.format(dataset))
                raise Exception('Missing dimensions Attr in dataset')
            else:
                dimensions = f['data'][dataset].attrs.get('dimensions')
                if isinstance(dimensions, bytes):
                    dimensions = dimensions.decode().split(',')
                else:
                    dimensions = dimensions.split(',')
                axis_data.append(dimensions)
                for aa in dimensions:
                    ls_axis.append(aa)
        else:
            print('dataset "{}" not available.'.format(dataset))
            print('{} contains the following datasets:'.format(filename))
            for d in f['data'].keys():
                print ('  *', d)
            raise Exception('Missing dataset')
    ls_axis = list(set(ls_axis))

    m = MLUT()
    # add axis to the MLUT
    for ax in ls_axis:
        axis = f['axis'][ax][...]
        m.add_axis(ax, axis)

    # add data to MLUT
    for idata in xrange(len(ls_datasets)):
        dataset = ls_datasets[idata]
        if lazy:
            data = f['data'][dataset]
        else:
            data = f['data'][dataset][...]

        if wrap_data is not None:
            data = wrap_data(data)

        attrs = {}
        if f['data'][dataset].attrs.__contains__('_FillValue'):
            attrs['_FillValue'] = f['data'][dataset].attrs.get('_FillValue')
        if f['data'][dataset].attrs.__contains__('add_offset'):
            attrs['add_offset'] = f['data'][dataset].attrs.get('add_offset')
        if f['data'][dataset].attrs.__contains__('scale_factor'):
            attrs['scale_factor'] = f['data'][dataset].attrs.get('scale_factor')
        m.add_dataset(dataset, data, axnames=axis_data[idata], attrs=attrs)

    if not lazy:
        ff.close()

    return m


def read_mlut_hdf(filename, datasets=None):
    '''
    read a MLUT from a hdf file (filename)
    datasets: list of datasets to read:
        * None (default): read all datasets, including axes as indicated by the
          attribute 'dimensions'
        * a list of:
            - dataset names (string)
            - or a tuple (dataset_name, axes) where axes is a list of
              dimensions (strings), overriding the attribute 'dimensions'
    '''
    from pyhdf.SD import SD

    hdf = SD(filename)

    # read the datasets
    ls_axes = []
    ls_datasets = []
    if datasets is None:
        datasets = xrange(len(hdf.datasets()))
    else:
        assert isinstance(datasets, list), 'datasets should be provided as a list'

    for i in datasets:
        if isinstance(i, tuple):
            (name, axes) = i
            sds = hdf.select(name)
        else:
            axes = None
            sds = hdf.select(i)
        sdsname = sds.info()[0]

        if (axes is None) and ('dimensions' in sds.attributes()):
            axes = sds.attributes()['dimensions'].split(',')
            axes = [x.strip() for x in axes]

            # replace 'None's by None
            axes = [None if (x=='None') else x for x in axes]

        if axes is not None:
            ls_axes.extend(axes)

        data = sds.get()
        attrs = sds.attributes()
        if 'lut:scalar' in attrs:
            attrs.pop('lut:scalar')
            data = data.reshape(())
        ls_datasets.append((sdsname, data, axes, attrs))

    # remove 'None' axes
    while None in ls_axes:
        ls_axes.remove(None)

    # transfer the axes from ls_datasets to the new MLUT
    m = MLUT()
    for ax in set(ls_axes):
        [x[0] for x in ls_datasets]

        # read the axis if not done already
        if ax not in [x[0] for x in ls_datasets]:
            if ax in hdf.datasets():
                sds = hdf.select(ax)
                m.add_axis(ax, sds.get())
        else:
            i = [x[0] for x in ls_datasets].index(ax)
            (name, data, _, _) = ls_datasets.pop(i)
            m.add_axis(name, data)

    # add the datasets
    for (name, data, axnames, attrs) in ls_datasets:
        if 'dimensions' in attrs:
            attrs.pop('dimensions')
        m.add_dataset(name, data, axnames, attrs)

    # read the global attributes
    for k, v in hdf.attributes().items():
        m.set_attr(k, v)

    return m

def from_xarray(A):
    """
    Convert xarray object `A` (`Dataset` or `DataArray`) to MLUT or LUT.
    """
    if isinstance(A, xr.Dataset):
        m = MLUT()

        for x in A:
            m.add_dataset(
                x,
                A[x].data,
                A[x].dims,
                attrs=A[x].attrs)
        
        for x in A.coords:
            m.add_axis(x, A.coords[x].data)

        return m

    elif isinstance(A, xr.DataArray):
        return LUT(
            A.data,
            axes=[A.coords[x].data for x in A.dims],
            names=A.dims,
            desc=A.name,
            attrs=A.attrs)

    else:
        raise Exception(f'Invalid class passed to `from_xarray` ({A.__class__})')



if __name__ == '__main__':
    import doctest
    doctest.testmod()
