# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2016-2017 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/


__authors__ = ["J. Garriga"]
__license__ = "MIT"
__date__ = "15/05/2020"

import copy
import glob
import logging
import os
import threading
from enum import IntEnum

import cv2
import h5py
import numpy
from sklearn import decomposition

from silx.io import utils
from silx.io import fabioh5
from silx.io.url import DataUrl

from darfix.core.imageOperations import img2img_mean, chunk_image
from darfix.core.imageOperations import background_subtraction, background_subtraction_2D
from darfix.core.imageOperations import hot_pixel_removal_3D, hot_pixel_removal_2D
from darfix.core.imageOperations import Method
from darfix.core.imageRegistration import shift_detection, apply_shift
from darfix.core.roi import apply_2D_ROI, apply_3D_ROI
from darfix.decomposition.ipca import IPCA
from darfix.decomposition.nmf import NMF
from darfix.decomposition.nica import NICA
from darfix.io import utils as io_utils

_logger = logging.getLogger(__file__)

DEFAULT_METADATA = fabioh5.FabioReader.DEFAULT

COUNTER_METADATA = fabioh5.FabioReader.COUNTER

POSITIONER_METADATA = fabioh5.FabioReader.POSITIONER

_METADATA_TYPES = {'default': DEFAULT_METADATA,
                   'counter': COUNTER_METADATA,
                   'positioner': POSITIONER_METADATA}

_METADATA_TYPES_I = {}
"""used to retrieve the metadata name (str) for the silx.io.fabioh5 id"""
for key, value in _METADATA_TYPES.items():
    assert value not in _METADATA_TYPES_I
    _METADATA_TYPES_I[value] = key


class Operation(IntEnum):
    """
    Flags for different operations in Dataset
    """
    PARTITION = 0
    BS = 1
    HP = 2
    SHIFT = 3


class Dataset():
    """Class to define a dataset from a series of data files.

    :param _dir: Global directory to use and save all the data in the different
        operations.
    :type _dir: str
    :param data: If not None, sets the Data array with the data of the dataset
        to use, default None
    :type data: :class:`Data`, optional
    :param raw_folder: Path to the folder that contains the data, defaults to
        None
    :type raw_folder: Union[None,str], optional
    :param filenames: Ordered list of filenames, defaults to None.
    :type filenames: Union[Generator,Iterator,List], optional
    :param treated: Name of the folder with the treated data, defaults 'treated_data'
    :type treated: str, optional
    :param in_memory: If True, data is loaded into memory, else, data is read in
        chunks depending on the algorithm to apply, defaults to False.
    :type in_memory: bool, optional
    """

    def __init__(self, _dir, data=None, raw_folder=None, filenames=None,
                 dims=None, treated="treated_data", in_memory=True):

        self._data = None
        self._frames_intensity = []
        self._dir = _dir
        self._treated = treated
        self._lock = threading.Lock()
        self.operations_state = numpy.zeros(len(Operation))
        if not os.path.isdir(self.dir + "/" + treated):
            os.mkdir(self.dir + "/" + treated)

        if dims is None:
            self.__dims = AcquisitionDims()
        else:
            assert isinstance(dims, AcquisitionDims), "Attribute dims has to be of class AcquisitionDims"
            self.__dims = dims
        # Keys: dimensions names, values: dimensions values
        self._dimensions_values = {}
        self._in_memory = in_memory

        if data is not None:
            self._data = data
        else:
            if filenames is None:
                if raw_folder is None:
                    filenames = sorted([x for x in glob.glob(self.dir + "/*") if os.path.isfile(x)])
                else:
                    filenames = sorted([x for x in glob.glob(self.dir + "/" + raw_folder + "/*") if os.path.isfile(x)])
            self.filenames = filenames

            metadata = []
            data_urls = []

            for filename in self.filenames:
                data_urls.append(DataUrl(file_path=filename,
                                         scheme='fabio'))
                fabio_reader = fabioh5.EdfFabioReader(file_name=filename)
                metadata.append(fabio_reader)
                fabio_reader.close()

            self._data = Data(numpy.array(data_urls), metadata=metadata, in_memory=self._in_memory)

    def stop_operation(self, operation):
        """
        Method used for cases where threads are created to apply functions to the dataset.
        If method is called, the flag concerning the stop is set to 0 so that if the concerned
        operation is running in another thread it knows to stop.

        :param int operation: operation to stop
        :type int: Union[int, `Operation`]
        """
        if self.operations_state[operation]:
            self._lock.acquire()
            self.operations_state[operation] = 0
            self._lock.release()
            self.get_data().stop_operation(operation)

    @property
    def dir(self):
        return self._dir

    @property
    def treated(self):
        return self._treated

    def compute_frames_intensity(self, kernel=(3, 3), sigma=20):
        """
        Returns for every image a number representing its intensity. This number
        is obtained by first blurring the image and then computing its variance.
        """
        _logger.info("Computing intensity per frame")
        io_utils.advancement_display(0, self.nframes, "Computing intensity")
        frames_intensity = []
        self._lock.acquire()
        self.operations_state[Operation.PARTITION] = 1
        self._lock.release()

        for i in range(self.nframes):
            if not self.operations_state[Operation.PARTITION]:
                return
            frames_intensity += \
                [cv2.GaussianBlur(self.get_data(i), kernel, sigma).var()]
            io_utils.advancement_display(i + 1, self.nframes, "Computing intensity")
        self._frames_intensity = frames_intensity
        self._lock.acquire()
        self.operations_state[Operation.PARTITION] = 0
        self._lock.release()
        return self._frames_intensity

    def partition_by_intensity(self, bins=None, num_bins=1):
        """
        Function that computes the data from the set of urls.
        If the filter_data flag is activated it filters the data following the next:
        -- First, it computes the intensity for each frame, by calculating the variance after
        passing a gaussian filter.
        -- Second, computes the histogram of the intensity.
        -- Finally, saves the data of the frames with an intensity bigger than a threshold.
        The threshold is set to be the second bin of the histogram.

        :param int num_bins: Number of bins to use as threshold.
        """
        frames_intensity = self._frames_intensity if self._frames_intensity else self.compute_frames_intensity()
        if frames_intensity is None:
            return
        values, bins = numpy.histogram(frames_intensity, self.nframes if bins is None else bins)
        threshold = numpy.array(frames_intensity) >= bins[num_bins]
        return numpy.flatnonzero(threshold), numpy.flatnonzero(~threshold)

    @property
    def in_memory(self):
        return self._in_memory

    @in_memory.setter
    def in_memory(self, in_memory):
        """
        Removes data from memory and sets that from now on data will be read from disk.
        """
        if self._in_memory is not in_memory:
            self._in_memory = in_memory
            self._data = Data(self.data.urls, self.data.metadata, self._in_memory)

    @property
    def data(self):
        return self._data

    def get_data(self, indices=None, dimension=None):
        """
        Returns the data corresponding to a certains indices and given a certain dimension.
        The data is always flattened to be a stack of images.

        :param array_like indices: If not None, data is filtered using this array.
        :param array_like dimension: If not None, return only the data corresponding to
            the given dimension. Dimension is a 2d vector, where the first component is
            the axis and the second is the indices of the values to extract.

        :return: Array with the new data.
        """
        if dimension is not None and len(self._data.shape) > 3:
            data = self.data.take(indices=dimension[1], axis=dimension[0])
            if indices is not None:
                bool_indices = numpy.zeros(self.data.flatten().shape[:-2], dtype=bool)
                bool_indices[indices] = True
                bool_indices = bool_indices.reshape(self.data.shape[:-2])
                indices = numpy.take(bool_indices, indices=dimension[1], axis=dimension[0])
                data = data[indices]
            return data
        else:
            data = self.data.flatten()
            return data if indices is None else data[indices]

    @property
    def nframes(self):
        """
        Return number of frames
        """
        if self.data is None:
            return 0
        else:
            return len(self.data.flatten())

    def to_memory(self, indices):
        if not self._in_memory:
            data = self.get_data(indices)
            new_data = Data(data.urls, data.metadata, True)
        else:
            new_data = self.get_data(indices)
        return Dataset(_dir=self._dir, data=new_data, dims=self.__dims, in_memory=True)

    @property
    def dims(self):
        return self.__dims

    @dims.setter
    def dims(self, _dims):
        assert isinstance(_dims, AcquisitionDims), "Dimensions dictionary has " \
            "to be of class `AcquisitionDims`"
        self.__dims = _dims

    def clear_dims(self):
        self.__dims = AcquisitionDims()

    def add_dim(self, axis, dim):
        """
        Adds a dimension to the dimension's dictionary.

        :param int axis: axis of the dimension.
        :param `Dimension` dim: dimension to be added.
        """
        self.__dims.add_dim(axis, dim)

    def remove_dim(self, axis):
        """
        Removes a dimension from the dimension's dictionary.

        :param int axis: axis of the dimension.
        """
        self.__dims.remove_dim(axis)

    def zsum(self, indices=None):

        if self._in_memory:
            return self.get_data(indices).sum(axis=0)
        else:
            if indices is None:
                indices = range(self.nframes)
            zsum = self.get_data(indices[0])
            for i in indices[1:]:
                zsum += self.get_data(i)
            return zsum

    def reshape_data(self):
        """
        Function that reshapes the data to fit the dimensions.
        """
        if self.__dims.ndim > 1:
            try:
                shape = list(self.__dims.shape)
                shape.append(self.get_data().shape[-2])
                shape.append(self.get_data().shape[-1])
                data = self.get_data().reshape(shape)
                return Dataset(_dir=self.dir, data=data, dims=self.__dims, in_memory=self._in_memory)
            except Exception:
                raise ValueError("Failed to reshape data into dimensions {} \
                                  Try using another tolerance value.".format(' '.join(self.__dims.get_names())))
        else:
            raise ValueError("Not enough dimensions where found")

    def find_dimensions(self, kind, tolerance=1e-9):
        """
        Goes over all the headers from a given kind and finds the dimensions
        that move (have more than one value) along the data.

        Note: Before, only the dimensions that could fit where shown, now it
        shows all the dimensions and let the user choose the valid ones.

        :param int kind: Type of metadata to find the dimensions.
        :param float tolerance: Tolerance that will be used to compute the
            unique values.
        """
        self.__dims.clear()
        self._dimensions_values = {}

        keys = numpy.array(list(self.data.metadata[0].get_keys(kind)))
        values = numpy.array([[data.get_value(kind=kind, name=key)[0] for data
                             in self.data.metadata] for key in keys])
        # Unique values for each key.
        unique_values = [numpy.unique(value, return_counts=True) for value in values]
        dimensions = []
        dataset_size = len(self.data.metadata)
        # For every key that has more than one different value, creates a new Dimension.
        for i, value in enumerate(unique_values):
            if value[1][0] != dataset_size:
                dimension = Dimension(kind, keys[i], tolerance=tolerance)
                dimension.set_unique_values(numpy.unique(value[0]))
                # Value that tells when does the change of value occur. It is used to know the order
                # of the reshaping.
                dimension.changing_value = numpy.unique(values[i, :int(dataset_size / value[1][0])],
                                                        return_counts=True)[1][0]
                dimensions.append(dimension)

        for dimension in sorted(dimensions, key=lambda x: x.changing_value):
            self.__dims.add_dim(axis=self.__dims.ndim, dim=dimension)
            _logger.info("Dimension {} of size {} has been added for reshaping"
                         .format(dimension.name, dimension.size))

    def get_dimensions_values(self, indices=None):
        """
        Returns all the metadata values of the dimensions.
        The values are assumed to be numbers.

        :returns: array_like
        """
        if not self._dimensions_values:
            data = self.get_data(indices).metadata
            for dimension in self.__dims:
                values = numpy.empty((len(data)))
                for row, metadata_frame in enumerate(data):
                    values[row] = (metadata_frame.get_value(kind=dimension[1].kind,
                                   name=dimension[1].name)[0])
                self._dimensions_values[dimension[1].name] = values
        return self._dimensions_values

    def apply_background_subtraction(self, background=None, method="median", indices=None,
                                     chunk_shape=[100, 100]):
        """
        Applies background subtraction to the data and saves the new data
        into disk.

        :param background: Data to be used as background. If None, data with indices `indices` is used.
            If Dataset, data of the dataset is used. If array, use data with indices in the array.
        :type background: Union[None, array_like, Dataset]
        :param method: Method to use to compute the background.
        :type method: Method
        :param indices: Indices of the images to apply background subtraction.
            If None, the background subtraction is applied to all the data.
        :type indices: Union[None, array_like]
        :param chunk_shape: Parameter used only when flag in_memory is False and method is Method.median.
            It is the shape of the chunk image to use per iteration.
        :type chunk_shape: array_like
        :returns: dataset with data of same size as `self.data` but with the
            modified images. The urls of the modified images are replaced with
            the new urls.
        :rtype: Dataset
        """
        bs_dir = self.dir + "/" + self.treated + "/bs"
        if not os.path.isdir(bs_dir):
            os.mkdir(bs_dir)

        data = self.get_data(indices)

        self._lock.acquire()
        self.operations_state[Operation.BS] = 1
        self._lock.release()

        if background is None:
            bg_data = data
            if indices is None:
                _logger.info("Computing background from " + method + " of raw data")
            else:
                _logger.info("Computing background from " + method + " of high intensity data")
        elif isinstance(background, Dataset):
            bg_data = background.data
            _logger.info("Computing background from " + method + " of `background` set")
        else:
            bg_data = self.get_data(background)
            _logger.info("Computing background from " + method + " of low intensity data")

        if self._in_memory:
            new_data = background_subtraction(data, bg_data, method).view(Data)
            new_data.save(bs_dir + "/data_bs")
            urls = new_data.urls
        else:
            method = Method.from_value(method)
            bg = numpy.zeros(data[0].shape, data.dtype)
            if method == Method.mean:
                if bg_data.in_memory:
                    numpy.mean(bg_data, out=bg, axis=0)
                else:
                    io_utils.advancement_display(0, len(bg_data), "Computing mean image")
                    for i in range(len(bg_data)):
                        if not self.operations_state[Operation.BS]:
                            return
                        bg = img2img_mean(bg_data[i], bg, i)
                        io_utils.advancement_display(i + 1, len(bg_data), "Computing mean image")
            elif method == Method.median:
                if bg_data.in_memory:
                    numpy.median(bg_data, out=bg, axis=0)
                else:
                    start = [0, 0]
                    img = data[0]
                    bg = numpy.empty(img.shape, img.dtype)
                    chunks_1 = int(numpy.ceil(img.shape[0] / chunk_shape[0]))
                    chunks_2 = int(numpy.ceil(img.shape[1] / chunk_shape[1]))
                    io_utils.advancement_display(0, chunks_1 * chunks_2 * len(bg_data), "Computing median image")
                    for i in range(chunks_1):
                        if not self.operations_state[Operation.BS]:
                            return
                        for j in range(chunks_2):
                            c_images = []
                            for k, img in enumerate(bg_data):
                                c_images += [chunk_image(img, start, chunk_shape)]
                                io_utils.advancement_display(
                                    i * chunks_2 * len(bg_data) + j * len(bg_data) + k + 1,
                                    chunks_1 * chunks_2 * len(bg_data),
                                    "Computing median image")
                            numpy.median(c_images, out=bg[start[0]: start[0] + chunk_shape[0],
                                         start[1]:start[1] + chunk_shape[1]], axis=0)
                            start[1] = chunk_shape[1] * (j + 1)
                        start[0] = chunk_shape[0] * (i + 1)
                        start[1] = 0
            urls = data.apply_funcs([(background_subtraction_2D, [bg])],
                                    save=bs_dir + "/data_bs", text="Applying background subtraction",
                                    operation=Operation.BS)
            if urls is None:
                return

        self._lock.acquire()
        self.operations_state[Operation.BS] = 0
        self._lock.release()

        # Set urls as shape and dimension of original urls.
        if indices is not None:
            new_urls = numpy.array(self.get_data().urls, dtype=object)
            new_urls[indices] = urls
            new_data = Data(new_urls.reshape(self.data.urls.shape), self.data.metadata,
                            self._in_memory)
        else:
            new_data = Data(urls.reshape(self.data.urls.shape), self.data.metadata,
                            self._in_memory)

        return Dataset(_dir=bs_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)

    def apply_hot_pixel_removal(self, kernel=3, indices=None):
        """
        Applies hot pixel removal to Data, and saves the new data
        into disk.

        :param int kernel: size of the kernel used to find the hot
            pixels.
        :param indices: Indices of the images to apply background subtraction.
            If None, the hot pixel removal is applied to all the data.
        :type indices: Union[None, array_like]
        :return: dataset with data of same size as `self.data` but with the
            modified images. The urls of the modified images are replaced with
            the new urls.
        :rtype: Dataset
        """

        data = self.get_data(indices)
        _dir = self.dir + "/" + self.treated + "/hp"
        if not os.path.isdir(_dir):
            os.mkdir(_dir)

        if self._in_memory:
            new_data = hot_pixel_removal_3D(data, kernel).view(Data)
            new_data.save(_dir + "/data_hp")
            urls = new_data.urls
        else:
            urls = data.apply_funcs([(hot_pixel_removal_2D, [kernel])],
                                    save=_dir + "/data_hp", text="Applying hot pixel removal",
                                    operation=Operation.HP)
            if urls is None:
                return

        if indices is not None:
            new_urls = numpy.array(self.get_data().urls, dtype=object)
            new_urls[indices] = urls
            new_data = Data(new_urls.reshape(self.data.urls.shape), self.data.metadata,
                            self._in_memory)
        else:
            new_data = Data(urls.reshape(self.data.urls.shape), self.data.metadata,
                            self._in_memory)

        return Dataset(_dir=_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)

    def apply_roi(self, origin=None, size=None, center=None, indices=None):
        """
        Applies a region of interest to the data.

        :param origin: Origin of the roi
        :param size: [Height, Width] of the roi.
        :param center: Center of the roi
        :type origin: Union[2d vector, None]
        :type center: Union[2d vector, None]
        :type size: Union[2d vector, None]
        :param indices: Indices of the images to apply background subtraction.
            If None, the roi is applied to all the data.
        :type indices: Union[None, array_like]
        :returns: dataset with data with roi applied.
            Note: To preserve consistence of shape between images, if `indices`
            is not None, only the data modified is returned.
        :rtype: Dataset
        """

        roi_dir = self.dir + "/" + self.treated + "/roi"
        if not os.path.isdir(roi_dir):
            os.mkdir(roi_dir)
        if self._in_memory:
            new_data = apply_3D_ROI(self.get_data(indices), origin, size, center).view(Data)
            new_data.save(roi_dir + "/roi")
        else:
            urls = self.get_data(indices).apply_funcs([(apply_2D_ROI, [origin, size, center])],
                                                      save=roi_dir + "/data_roi", text="Applying roi")
            if urls is None:
                return
            new_data = Data(urls, self.get_data(indices).metadata, self._in_memory)
        if indices is None:
            shape = list(self.data.shape)[:-2]
            shape.append(new_data.shape[-2])
            shape.append(new_data.shape[-1])
            new_data = new_data.reshape(shape)
        return Dataset(_dir=roi_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)

    def find_shift(self, dimension=None, h_max=0.5, h_step=0.01, indices=None):
        """
        Find shift of the data or part of it.

        :param dimension: Parametes with the position of the data in the reshaped
            array.
        :type dimension: Union[None, tuple, array_like]
        :param float h_max: See `core.imageRegistration.shift_detection`
        :param float h_step: See `core.imageRegistration.shift_detection`
        :param indices: Boolean index list with True in the images to apply the shift to.
            If None, the hot pixel removal is applied to all the data.
        :type indices: Union[None, array_like]
        :returns: Array with shift per frame.
        """
        return shift_detection(self.get_data(indices, dimension), h_max, h_step)

    def apply_shift(self, shift, dimension=None, shift_approach="fft", indices=None, callback=None):
        """
        Apply shift of the data or part of it and save new data into disk.

        :param array_like shift: Shift per frame.
        :param dimension: Parametes with the position of the data in the reshaped
            array.
        :type dimension: Union[None, tuple, array_like]
        :param Union['fft', 'linear'] shift_approach: Method to use to apply the shift.
        :param indices: Boolean index list with True in the images to apply the shift to.
            If None, the hot pixel removal is applied to all the data.
        :type indices: Union[None, array_like]
        :param Union[function, None] callback: Callback
        :returns: dataset with data of same size as `self.data` but with the
            modified images. The urls of the modified images are replaced with
            the new urls.
        """
        assert len(shift) > 0, "Shift list can't be empty"

        shift_dir = self.dir + "/" + self.treated + "/shift"
        if not os.path.isdir(shift_dir):
            os.mkdir(shift_dir)

        data = self.get_data(indices, dimension)
        self._lock.acquire()
        self.operations_state[Operation.SHIFT] = 1
        self._lock.release()

        io_utils.advancement_display(0, len(data), "Applying shift")
        if dimension is not None:
            urls = []
            for i in range(len(data)):
                if not self.operations_state[Operation.SHIFT]:
                    return
                filename = shift_dir + "/data_shift" + str(i).zfill(4) + ".npy"
                img = apply_shift(data[i], shift[:, i], shift_approach)
                numpy.save(filename, img)
                urls.append(DataUrl(file_path=filename, scheme='fabio'))
                io_utils.advancement_display(i + 1, len(data), "Applying shift")
            new_urls = numpy.array(self.data.urls, dtype=object)
            if indices is not None:
                # Create array of booleans to know which indices we have
                bool_indices = numpy.zeros(self.get_data().shape[:-2], dtype=bool)
                bool_indices[indices] = True
                bool_indices = bool_indices.reshape(self.data.shape[:-2])
                ind = numpy.swapaxes(bool_indices, 0, dimension[0])[dimension[1], :]
                numpy.swapaxes(new_urls, 0, dimension[0])[dimension[1], :][ind] = urls
            else:
                numpy.swapaxes(new_urls, 0, dimension[0])[dimension[1], :] = urls
        else:
            urls = []
            for i in range(len(data)):
                if not self.operations_state[Operation.SHIFT]:
                    return
                filename = shift_dir + "/data_shift" + str(i).zfill(4) + ".npy"
                img = apply_shift(data[i], shift[:, i], shift_approach)
                numpy.save(filename, img)
                urls.append(DataUrl(file_path=filename, scheme='fabio'))
                io_utils.advancement_display(i + 1, len(data), "Applying shift")
            if indices is not None:
                new_urls = numpy.array(self.data.urls, dtype=object).flatten()
                numpy.put(new_urls, indices, urls)
            else:
                new_urls = numpy.array(urls)

        self._lock.acquire()
        self.operations_state[Operation.SHIFT] = 0
        self._lock.release()

        data = Data(new_urls.reshape(self.data.urls.shape), self.data.metadata, in_memory=self._in_memory)
        return Dataset(_dir=shift_dir, data=data, dims=self.__dims, in_memory=self._in_memory)

    def find_and_apply_shift(self, dimension=None, h_max=0.5, h_step=0.01, shift_approach="fft",
                             indices=None, callback=None):
        """
        Find the shift of the data or part of it and apply it.

        :param dimension: Parametes with the position of the data in the reshaped
            array.
        :type dimension: Union[None, tuple, array_like]
        :param float h_max: See `core.imageRegistration.shift_detection`
        :param float h_step: See `core.imageRegistration.shift_detection`
        :param Union['fft', 'linear'] shift_approach: Method to use to apply the shift.
        :param indices: Indices of the images to find and apply the shift to.
            If None, the hot pixel removal is applied to all the data.
        :type indices: Union[None, array_like]
        :param Union[function, None] callback: Callback
        :returns: Dataset with the new data.
        """
        shift = self.find_shift(dimension, h_max, h_step, indices=indices)
        return self.apply_shift(shift, dimension, indices=indices)

    def _cascade_nmf(self, num_components, iterations, vstep=None, hstep=None, indices=None):
        """
        This method is used as a way to improve the speed of convergence of
        the NMF method. For this, it uses a cascade methology at every
        step of the cascade, the data used has a smaller size.
        That is, the method starts with a smaller resized images of the data,
        and computes the NMF decomposition. The next step is the same but with
        bigger images, and using as initial H and W the precomputed matrices.
        The last step is done using the actual size of the images.
        This way, the number of iterations with big images can be diminished, and
        the method converges faster.

        :param int num_components: Number of components to find.
        :param array_like iterations: Array with number of iterations per step of the cascade.
            The size of the array sets the size of the cascade.
        :param Union[None, array_like] indices: If not None, apply method only to indices of data.
        """

        from skimage.transform import resize
        import shutil

        W = None
        H = None
        shape = numpy.asarray(self.get_data(0).shape)
        first_size = (shape / (len(iterations) + 1)).astype(numpy.int)
        size = first_size

        if not os.path.isdir(self.dir + "/cascade"):
            os.mkdir(self.dir + "/cascade")

        _logger.info("Starting cascade NMF")

        for i in range(len(iterations)):
            new_urls = []
            if indices is None:
                indices = range(self.nframes)
            for j in indices:
                filename = self.dir + "/cascade/" + str(j) + ".npy"
                numpy.save(filename, resize(self.get_data(j), size))
                new_urls.append(DataUrl(file_path=filename, scheme='fabio'))
            data = Data(new_urls, self.get_data(indices).metadata, self._in_memory)
            dataset = Dataset(_dir=self.dir + "/cascade", data=data, in_memory=self._in_memory)
            if vstep:
                v_step = vstep * (len(iterations) - i)
            if hstep:
                h_step = hstep * (len(iterations) - i)
            H, W = dataset.nmf(num_components, iterations[i], W=W, H=H, vstep=v_step, hstep=h_step)
            size = first_size * (i + 2)
            H2 = numpy.empty((H.shape[0], size[0] * size[1]))
            for row in range(H.shape[0]):
                H2[row] = resize(H[row].reshape((i + 1) * first_size), size).flatten()
            H = H2

        try:
            shutil.rmtree(self.dir + "/cascade")
        except Exception as e:
            print('Failed to delete %s. Reason: %s' % (self.dir + "/cascade", e))

        H = resize(H, (num_components, shape[0] * shape[1]))

        return H, W

    def pca(self, num_components=None, chunk_size=None, indices=None, return_vals=False):
        """
        Compute Principal Component Analysis on the data.
        The method, first converts, if not already, the data into an hdf5 file object
        with the images flattened in the rows.

        :param num_components: Number of components to find.
            If None, it uses the minimum between the number of images and the
            number of pixels.
        :type num_components: Union[None, int]
        :param chunk_size: Number of chunks for which the whitening must be computed,
            incrementally, defaults to None
        :type chunksize: Union[None, int], optional
        :param indices: If not None, apply method only to indices of data, defaults to None
        :type indices: Union[None, array_like], optional
        :param return_vals: If True, returns only the singular values of PCA, else returns
            the components and the mixing matrix, defaults to False
        :type return_vals: bool, optional

        :return: (H, W): The components matrix and the mixing matrix.
        """
        bss_dir = self.dir + "/" + self.treated + "/bss"
        if not os.path.isdir(bss_dir):
            os.mkdir(bss_dir)
        if self._in_memory:
            model = decomposition.PCA(n_components=num_components)
            try:
                if indices is not None:
                    W = model.fit_transform(self.data.convert_to_hdf5(bss_dir)[indices])
                else:
                    W = model.fit_transform(self.data.convert_to_hdf5(bss_dir)[:, :])
            finally:
                self.data._file.close()

            H, vals, W = model.components_, model.singular_values_, W
        else:
            if chunk_size is None:
                chunk_size = int(self.get_data().shape[0] / 10)
            model = IPCA(self.data.convert_to_hdf5(bss_dir), chunk_size, num_components, indices=indices)
            file = h5py.File(bss_dir + '/ipca.hdf5', 'w')
            H, W, vals = None, None, None
            try:
                file["W"] = numpy.random.random((model.num_samples, model.num_components))
                file["H"] = numpy.random.random((model.num_components, model.num_features))
                model.fit_transform(W=file["W"], H=file["H"])
            finally:
                file.close()
                self.data._file.close()

            H, vals, W = model.H, model.singular_values, model.W
        return vals if return_vals else (H, W)

    def nica(self, num_components, chunksize=None, num_iter=500, error_step=None, indices=None):
        """
        Compute Non-negative Independent Component Analysis on the data.
        The method, first converts, if not already, the data into an hdf5 file object
        with the images flattened in the rows.

        :param num_components: Number of components to find
        :type num_components: Union[None, int]
        :param chunksize: Number of chunks for which the whitening must be computed,
            incrementally, defaults to None
        :type chunksize: Union[None, int], optional
        :param num_iter: Number of iterations, defaults to 500
        :type num_iter: int, optional
        :param error_step: If not None, find the error every error_step and compares it
            to check for convergence. TODO: not able for huge datasets.
        :param indices: If not None, apply method only to indices of data, defaults to None
        :type indices: Union[None, array_like], optional

        :return: (H, W): The components matrix and the mixing matrix.
        """
        bss_dir = self.dir + "/" + self.treated + "/bss"
        if not os.path.isdir(bss_dir):
            os.mkdir(bss_dir)

        if self._in_memory:
            chunksize = None
        model = NICA(self.data.convert_to_hdf5(bss_dir), num_components, chunksize, indices=indices)
        try:
            model.fit_transform(max_iter=num_iter, error_step=error_step)
            return model.H, model.W
        finally:
            self.data._file.close()

    def nmf(self, num_components, num_iter=100, error_step=None, cascade=None,
            H=None, W=None, vstep=100, hstep=1000, indices=None):
        """
        Compute Non-negative Matrix Factorization on the data.
        The method, first converts, if not already, the data into an hdf5 file object
        with the images flattened in the rows.

        :param num_components: Number of components to find
        :type num_components: Union[None, int]
        :param num_iter: Number of iterations, defaults to 100
        :type num_iter: int, optional
        :param error_step: If not None, find the error every error_step and compares it
            to check for convergence, defaults to None
            TODO: not able for huge datasets.
        :type error_step: Union[None, int], optional
        :param cascade: If not None, NMF is computed using the cascade method.
            The parameter should be an array with the number of iterations per
            sub-computation, defaults to None
        :type cascade: Union[None, array_like], optional
        :param H: Init matrix for H, defaults to None
        :type H: Union[None, array_like], optional
        :param W: Init matrix for W, defaults to None
        :type W: Union[None, array_like], optional
        :param indices: If not None, apply method only to indices of data, defaults to None
        :type indices: Union[None, array_like], optional

        :return: (H, W): The components matrix and the mixing matrix.
        """
        bss_dir = self.dir + "/" + self.treated + "/bss"
        if not os.path.isdir(bss_dir):
            os.mkdir(bss_dir)

        if self._in_memory:
            model = decomposition.NMF(n_components=num_components, init='random', max_iter=num_iter)
            if indices is not None:
                X = self.data.convert_to_hdf5(bss_dir)[indices]
            else:
                X = self.data.convert_to_hdf5(bss_dir)
            if numpy.any(X[:, :] < 0):
                _logger.warning("Setting negative values to 0 to compute NMF")
                X[X[:, :] < 0] = 0
            try:
                W = model.fit_transform(X, W=W, H=H)
                return model.components_, W
            finally:
                self.data._file.close()
        else:
            if cascade is not None:
                H, W = self._cascade_nmf(num_components, cascade, vstep, hstep, indices=indices)

            model = NMF(self.data.convert_to_hdf5(bss_dir), num_components, indices=indices)
            try:
                model.fit_transform(max_iter=num_iter, H=H, W=W, vstep=vstep, hstep=hstep)
                return model.H, model.W
            finally:
                self.data._file.close()

    def nica_nmf(self, num_components, chunksize=None, num_iter=500, cascade=None,
                 error_step=None, vstep=100, hstep=1000, indices=None):
        """
        Applies both NICA and NMF to the data. The init H and W for NMF are the
        result of NICA.
        """
        H, W = self.nica(num_components, chunksize, num_iter, error_step, indices=indices)

        return self.nmf(num_components, num_iter, error_step, cascade, H, W,
                        vstep, hstep, indices=indices)

    def __deepcopy__(self, memo):
        """
        Create copy of the dataset. The data numpy array is also copied using
        deep copy. The rest of the attributes are the same.
        """
        dataset = type(self)(self.dir, data=self.data, filenames=self.filenames,
                             treated=self.treated)
        dataset.dims = copy.deepcopy(self.__dims, memo)
        return dataset


class Data(numpy.ndarray):
    """
    Class to structure the data and link every image with its corresponding url and
    metadata.

    :param urls: Array with the urls of the data
    :type urls: array_like
    :param metadata: Array with the metadata of the data
    :type metadata: array_like
    :param in_memory: If True, the data is loaded into memory, default True
    :type in_memory: bool, optional
    """
    def __new__(cls, urls, metadata, in_memory=True, data=None):
        urls = numpy.asarray(urls)
        if in_memory:
            if data is not None and urls.shape == data.shape[:-2]:
                input_data = data
            else:
                # Create array as stack of images
                input_data = []
                for url in urls.flatten():
                    input_data += [utils.get_data(url)]
                input_data = numpy.asarray(input_data)
            shape = list(urls.shape)
            shape.append(input_data.shape[-2])
            shape.append(input_data.shape[-1])
            obj = input_data.reshape(shape).view(cls)
        else:
            # Access image one at a time using url
            obj = super(Data, cls).__new__(cls, urls.shape)

        obj.in_memory = in_memory
        obj.urls = urls
        obj.metadata = numpy.asarray(metadata)
        obj._file = None

        obj._lock = threading.Lock()
        obj.operations = numpy.zeros(len(Operation))

        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return
        self.urls = getattr(obj, 'urls', None)
        self.metadata = getattr(obj, 'metadata', None)
        self.in_memory = getattr(obj, 'in_memory', None)

    def __getitem__(self, indices):
        """
        Return self[indices]
        """
        if self.in_memory:
            data = super(Data, self).__getitem__(indices)
            if len(data.shape) < 3:
                data = data.view(numpy.ndarray)
            else:
                if isinstance(indices[0], slice):
                    data.urls = self.urls[indices[0]]
                    data.metadata = self.metadata[indices[0]]
                else:
                    data.urls = self.urls[indices]
                    data.metadata = self.metadata[indices]
            return data

        if isinstance(indices, tuple):
            if not isinstance(self.urls[indices[0]], numpy.ndarray):
                return utils.get_data(self.urls[indices[0]])[indices[1], indices[2]]
            else:
                return Data(self.urls[indices[0]], self.metadata[indices], self.in_memory)
        else:
            if not isinstance(self.urls[indices], numpy.ndarray):
                return utils.get_data(self.urls[indices])
            else:
                return Data(self.urls[indices], self.metadata[indices], self.in_memory)

    def stop_operation(self, operation):
        """
        Method used for cases where threads are created to apply functions to the data.
        If method is called, the flag concerning the stop is set to 0 so that if the concerned
        operation is running in another thread it knows to stop.

        :param int operation: operation to stop
        :type int: Union[int, `Operation`]
        """
        if hasattr(self, 'operations') and self.operations[operation]:
            self._lock.acquire()
            self.operations[operation] = 0
            self._lock.release()

    @property
    def shape(self):
        """
        Tuple of array dimensions.
        """
        if self.in_memory:
            return super(Data, self).shape
        else:
            data = self.flatten()[0]
            shape = list(super(Data, self).shape)
            shape.append(data.shape[0])
            shape.append(data.shape[1])
            return tuple(shape)

    @property
    def ndim(self):
        if self.in_memory:
            return super(Data, self).ndim
        else:
            return super(Data, self).ndim + 2

    def apply_funcs(self, funcs=[], indices=None, save=False, text="", operation=None):
        """
        Method that applies a series of functions into the data. It can save the images
        into disk or return them.

        :param funcs: List of tupples. Every tupples contains the function to
            apply and its parameters, defaults to []
        :type funcs: array_like, optional
        :param indices: Indices of the data to apply the functions to,
            defaults to None
        :type indices: Union[None, array_like], optional
        :param save: If True, saves the images into disk, defaults to False
        :type save: bool
        :param str text: Text to show in the advancement display.
        :param int operation: operation to stop
        :type int: Union[int, `Operation`]

        :returns: Array with the new urls (if data was saved)
        """
        if indices is None:
            indices = range(len(self))
        if isinstance(indices, int):
            indices = [indices]
        urls = []
        io_utils.advancement_display(0, len(self.urls.flatten()), text)

        if not hasattr(self, 'operations'):
            self.operations = numpy.zeros(len(Operation))
            self._lock = threading.Lock()
        self._lock.acquire()
        self.operations[operation] = 1
        self._lock.release()
        for i in indices:
            if operation is not None and not self.operations[operation]:
                if save:
                    for j in indices:
                        while j != i:
                            filename = save + str(j).zfill(4) + ".npy"
                            os.remove(filename)
                return
            img = self[int(i)]
            for f, args in funcs:
                img = f(*([img] + args))
            if save:
                filename = save + str(i).zfill(4) + ".npy"
                numpy.save(filename, img)
                urls.append(DataUrl(file_path=filename, scheme='fabio'))
            io_utils.advancement_display(i + 1, len(self.urls.flatten()), text)
        self._lock.acquire()
        self.operations[operation] = 0
        self._lock.release()
        return numpy.array(urls)

    def save(self, path, indices=None):
        """
        Save the data into `path` folder and replace Data urls.
        TODO: check if urls already exist and, if so, modify only urls[indices].

        :param path: Path to the folder
        :type path: str
        :param indices: the indices of the values to save, defaults to None
        :type indices: Union[None,array_like], optional
        """
        self.in_memory = True
        urls = []
        if indices is None:
            data = self
        else:
            data = self[indices]
        for i, img in enumerate(data):
            filename = path + str(i).zfill(4) + ".npy"
            numpy.save(filename, img)
            urls.append(DataUrl(file_path=filename, scheme='fabio'))
        self.urls = numpy.asarray(urls)

    def convert_to_hdf5(self, _dir):
        """
        Converts the data into an HDF5 file, setting flattened images in the rows.
        TODO: pass filename per parameter?

        :param _dir: Directory in which to save the HDF5 file.
        :type _dir: str

        :return: Hdf5 file
        :rtype: `h5py.File`
        """

        self._file = h5py.File(_dir + '/data.hdf5', 'a')
        data = self.flatten()
        shape = data[0].shape[0] * data[0].shape[1]
        if "dataset" in self._file:
            if self._file["dataset"].shape == (len(data), shape):
                return self._file["dataset"]
            else:
                del self._file["dataset"]
        self._file.create_dataset("dataset", (len(data), shape))
        for i in range(len(data)):
            self._file["dataset"][i] = data[i].flatten()

        return self._file["dataset"]

    def reshape(self, shape):
        """
        Returns an array containing the same data with a new shape of urls and metadata.
        Shape also contains image shape at the last two positions (unreshapable).

        :param shape: New shape, should be compatible with the original shape.
        :type shape: int or tuple of ints.
        :return: new Data object with urls and metadata reshaped to shape.
        """
        data = None
        if self.in_memory:
            data = super(Data, self).reshape(shape).view(numpy.ndarray)
        return Data(self.urls.reshape(shape[:-2]), self.metadata.reshape(shape[:-2]),
                    self.in_memory, data=data)

    def take(self, indices, axis=None):
        """
        Take elements from urls and metadata from an array along an axis.

        :param indices: the indices of the values to extract
        :type indices: array_like
        :param axis: the axis over which to select values, defaults to None
        :type axis: Union[N one, int], optional

        :return: Flattened data.
        :rtype: :class:`Data`
        """
        urls = numpy.take(self.urls, indices, axis)
        metadata = numpy.take(self.urls, indices, axis)
        data = None
        if self.in_memory:
            data = numpy.take(self, indices, axis).view(numpy.ndarray)
        return Data(urls, metadata, self.in_memory, data=data)

    def flatten(self):
        """
        :return: new data with flattened urls and metadata (but not images).
        :rtype: :class:`Data`
        """
        data = None
        urls = self.urls.flatten()
        if self.in_memory:
            data = super(Data, self).reshape((len(urls), self.shape[-2], self.shape[-1])).view(numpy.ndarray)
        return Data(urls, self.metadata.flatten(), self.in_memory, data=data)


class AcquisitionDims(object):
    """
    Define the view of the data which has to be made
    """
    def __init__(self):
        self.__dims = {}

    def add_dim(self, axis, dim):
        assert isinstance(dim, Dimension)
        self.__dims[axis] = dim

    def remove_dim(self, axis):
        if axis in self.__dims:
            del self.__dims[axis]

    def clear(self):
        self.__dims = {}

    @property
    def ndim(self):
        return len(self.__dims)

    def get(self, axis):
        """
        Get Dimension at certain axis.

        :param int axis: axis of the dimension.
        :return: the requested dimension if exists.
        """
        assert type(axis) is int
        if axis in self.__dims:
            return self.__dims[axis]
        else:
            return None

    def get_names(self):
        """
        Get list with all the names of the dimensions.

        :return: array_like of strings
        """
        dims = []
        for dim in self.__dims.values():
            dims += [dim.name]

        return dims

    @property
    def shape(self):
        """
        :return: shape of the currently defined dims
        """
        shape = []
        for iDim in range(self.ndim):
            if iDim not in self.__dims:
                shape.append(1)
            else:
                shape.append(self.__dims[iDim].size or -1)
        return tuple(shape)

    def set_size(self, axis, size):
        """
        Recreated new dimension with new size and same name and kind.

        :param int axis: axis of the dimension
        :param int size: new size for the dimension
        """
        if axis not in self.__dims:
            _logger.error('axis %s is not defined yet, cannot define a size '
                          'for it' % axis)
        else:
            self.__dims[axis] = Dimension(name=self.__dims[axis].name,
                                          kind=self.__dims[axis].kind,
                                          size=size)

    def to_dict(self):
        dims = {}
        if self.ndim == 0:
            return
        else:
            for axis, dim in self.__dims.items():
                assert isinstance(dim, Dimension)
                dims[axis] = dim.to_dict()
            return dims

    def from_dict(self, _dict):
        """
        Convert dictionary of dimensions.
        """
        if len(_dict) > 0:
            for _axis, _dim in _dict.items():
                assert type(_dim) is dict
                self.add_dim(_axis, Dimension.from_dict(_dim))

    def __iter__(self):
        for iAxis, dim in self.__dims.items():
            yield (iAxis, dim)


class Dimension(object):
    """
    Define a dimension used during the dataset

    :param Union[int,str] kind: metadata type in fabioh5 mapping
    :param str name: name of the dimension (should fit the fabioh5 mapping
                     for now)
    :param Union[int,None] size: length of the dimension.
    """
    def __init__(self, kind, name, size=None, tolerance=1e-09):
        if type(kind) is str:
            assert kind in _METADATA_TYPES
            self.__kind = _METADATA_TYPES[kind]
        else:
            self.__kind = kind
        self.__name = name
        self._size = size
        self._tolerance = tolerance
        self.__unique_values = []

    @property
    def kind(self):
        return self.__kind

    def set_kind(self, kind):
        self.__kind = kind

    @property
    def name(self):
        return self.__name

    def set_name(self, name):
        self.__name = name

    @property
    def size(self):
        return self._size

    def set_size(self, size):
        self._size = size

    @property
    def tolerance(self):
        return self._tolerance

    def set_tolerance(self, tolerance):
        assert isinstance(tolerance, float), "Tolerance has to be float number"
        self._tolerance = tolerance

    @property
    def unique_values(self):
        return self.__unique_values

    def _find_unique_values(self, values):
        """
        Function that compares the values passed as parameter and returns only the unique
        ones given the dimension's tolerance.

        :param array_like values: list of values to compare.
        """
        unique_values = []
        import math
        if not numpy.all(numpy.isreal(values)):
            return values
        for val in values:
            if not unique_values:
                unique_values.append(val)
            else:
                unique = True
                for unique_value in unique_values:
                    if math.isclose(unique_value, val, rel_tol=self.tolerance):
                        unique = False
                        break
                if unique:
                    unique_values.append(val)
        return unique_values

    def set_unique_values(self, values):
        """
        Sets the unique values of the dimension. If the size of the dimension is fixed,
        it automatically sets the first size values, else it finds the unique values.

        :param array_like values: list of values.
        """
        if self.size:
            self.__unique_values = values[:self.size]
        else:
            self.__unique_values = self._find_unique_values(values)
        self.set_size(len(self.__unique_values))

    def __str__(self):
        return " ".join((str(self.kind), str(self.name), 'size:', str(self.size)))

    def to_dict(self):
        """Translate the current Dimension to a dictionary"""
        return {
            'name': self.name,
            'kind': self.kind,
            'size': self.size,
            'tolerance': self.tolerance
        }

    @staticmethod
    def from_dict(_dict):
        """

        :param dict _dict: dict defining the dimension. Should contains the
                           following keys: name, kind, size.
                           Unique values are not stored into it because it
                           depends on the metadata and should be obtained from a
                           fit / set_dims
        :return: Dimension corresponding to the dict given
        :rtype: :class:`Dimension`
        """
        assert type(_dict) is dict
        missing_keys = []
        for _key in ('name', 'kind', 'size', 'tolerance'):
            if _key not in _dict:
                missing_keys.append(missing_keys)
        if len(missing_keys) > 0:
            raise ValueError('There is some missing key (%s), unable to create'
                             'a valid Dim')
        else:
            return Dimension(name=_dict['name'],
                             kind=_dict['kind'],
                             size=_dict['size'],
                             tolerance=_dict['tolerance'])
