# -*- coding: utf-8 -*-
# Copyright 2007-2016 The HyperSpy developers
#
# This file is part of  HyperSpy.
#
#  HyperSpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
#  HyperSpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with  HyperSpy.  If not, see <http://www.gnu.org/licenses/>.

import copy
import math

import numpy as np
import dask.array as da
import traits.api as t
from traits.trait_errors import TraitError
import pint
import logging

from hyperspy.events import Events, Event
from hyperspy.misc.utils import isiterable, ordinal
from hyperspy.misc.math_tools import isfloat
from hyperspy.ui_registry import add_gui_method, get_gui
from hyperspy.defaults_parser import preferences


import warnings

_logger = logging.getLogger(__name__)
_ureg = pint.UnitRegistry()


FACTOR_DOCSTRING = \
    """factor : float (default: 0.25)
            'factor' is an adjustable value used to determine the prefix of
            the units. The product `factor * scale * size` is passed to the
            pint `to_compact` method to determine the prefix."""


class ndindex_nat(np.ndindex):

    def __next__(self):
        return super(ndindex_nat, self).next()[::-1]


def generate_axis(offset, scale, size, offset_index=0):
    """Creates an axis given the offset, scale and number of channels

    Alternatively, the offset_index of the offset channel can be specified.

    Parameters
    ----------
    offset : float
    scale : float
    size : number of channels
    offset_index : int
        offset_index number of the offset

    Returns
    -------
    Numpy array

    """

    return np.linspace(offset - offset_index * scale,
                       offset + scale * (size - 1 - offset_index),
                       size)


class UnitConversion:

    def __init__(self, units=t.Undefined, scale=1.0, offset=0.0):
        self.units = units
        self.scale = scale
   
     self.offset = units

    def _ignore_conversion(self, units):
        if units == t.Undefined:
            return True
        try:
            _ureg(units)
        except pint.errors.UndefinedUnitError:
            warnings.warn('Unit "{}" not supported for conversion. Nothing '
                          'done.'.format(units),
                          UserWarning)
            return True
        return False

    def _convert_compact_units(self, factor=0.25, inplace=True):
    
    """ Convert units to "human-readable" units, which means with a
            convenient prefix.

            Parameters
            ----------
            %s
        """
        if self._ignore_conversion(self.units):
            return
        scale = self.scale * _ureg(self.units)
        scale_size = factor * scale * self.size
        converted_units = '{:~}'.format(scale_size.to_compact().units)
        return self._convert_units(converted_units, inplace=inplace)

    _convert_compact_units.__doc__ %= FACTOR_DOCSTRING

    def _get_index_from_value_with_units(self, value):
        value = _ureg.parse_expression(value)
        if not hasattr(value, 'units'):
            raise ValueError('"{}" should contains an units.'.format(value))
        return self.value2index(value.to(self.units).magnitude)

    def _convert_units(self, converted_units, inplace=True):
        if self._ignore_conversion(converted_units) or \
                self._ignore_conversion(self.units):
            return
        scale_pint = self.scale * _ureg(self.units)
        offset_pint = self.offset * _ureg(self.units)
        scale = float(scale_pint.to(_ureg(converted_units)).magnitude)
        offset = float(offset_pint.to(_ureg(converted_units)).magnitude)
        units = '{:~}'.format(scale_pint.to(_ureg(converted_units)).units)
        if inplace:
            self.scale = scale
            self.offset = offset
            self.units = units
        else:
            return scale, offset, units

    def convert_to_units(self, units=None, inplace=True, factor=0.25):
        """ Convert the scale and the units of the current axis. If the unit
        of measure is not supported by the pint library, the scale and units
        are not modified.

        Parameters
        ----------
        units : {str | None}
            Default = None
            If str, the axis will be converted to the provided units.
            If `"auto"`, automatically determine the optimal units to avoid
            using too large or too small numbers. This can be tweaked by the
            `factor` argument.
        inplace : bool
            If `True`, convert the axis in place. if `False` return the
            `scale`, `offset` and `units`.
        %s
        """
        if units is None:
            out = self._convert_compact_units(factor, inplace=inplace)
        else:
            out = self._convert_units(units, inplace=inplace)
        return out

    convert_to_units.__doc__ %= FACTOR_DOCSTRING

    def _get_quantity(self, attribute='scale'):
        if attribute == 'scale' or attribute == 'offset':
            units = self.units
            if units == t.Undefined:
                units = ''
            return getattr(self, attribute) * _ureg(units)
        else:
            raise ValueError('`attribute` argument can only take the `scale` '
                             'or the `offset` value.')

    def _set_quantity(self, value, attribute='scale'):
        if attribute == 'scale' or attribute == 'offset':
            units = '' if self.units == t.Undefined else self.units
            if isinstance(value, str):
                value = _ureg.parse_expression(value)
            if isinstance(value, float):
                value = value * _ureg(units)

            # to be consistent, we also need to convert the other one
            # (scale or offset) when both units differ.
            if value.units != units and value.units != '' and units != '':
                other = 'offset' if attribute == 'scale' else 'scale'
                other_quantity = self._get_quantity(other).to(value.units)
                setattr(self, other, float(other_quantity.magnitude))

            self.units = '{:~}'.format(value.units)
            setattr(self, attribute, float(value.magnitude))
        else:
            raise ValueError('`attribute` argument can only take the `scale` '
                             'or the `offset` value.')

    @property
    def units(self):
        return self._units

    @units.setter
    def units(self, s):
        if s == '':
            self._units = t.Undefined
        self._units = s


@add_gui_method(toolkey="hyperspy.DataAxis")
class DataAxis(t.HasTraits, UnitConversion):
    name = t.Str()
    units = t.Str()
    scale = t.Float()
    offset = t.Float()
    size = t.CInt()
    low_value = t.Float()
    high_value = t.Float()
    value = t.Range('low_value', 'high_value')
    low_index = t.Int(0)
    high_index = t.Int()
    slice = t.Instance(slice)
    navigate = t.Bool(t.Undefined)
    index = t.Range('low_index', 'high_index')
    axis = t.Array()
    continuous_value = t.Bool(False)

    def __init__(self,
                 size,
                 index_in_array=None,
                 name=t.Undefined,
                 scale=1.,
                 offset=0.,
                 units=t.Undefined,
                 navigate=t.Undefined):
        super().__init__()
        self.events = Events()
        self.events.index_changed = Event("""
            Event that triggers when the index of the `DataAxis` changes

            Triggers after the internal state of the `DataAxis` has been
            updated.

            Arguments:
            ---------
            obj : The DataAxis that the event belongs to.
            index : The new index
            """, arguments=["obj", 'index'