"""A module for extending the standard math library."""

__author__ = "Damien RABOIS"


import math
# import bisect
import array
import operator
import copy
from Statistics import median


# Dependencies from CASSIS
try:
    import ScriptEnvironment
    from Plot import Plot
except ImportError:
    pass


# DECORATORS{{{1
def elementwise(fn):
    def newfn(*args):
        if hasattr(args[0], '__getitem__'):  
            return type(args[0])(map(lambda x: fn(x, *args[1:]), args[0]))
        else:
            return fn(*args)
    return newfn


def arenumbers(func):
    def wrapper(sample):
        if isiterable(sample):
            return func(filter(isnumber, sample))
        else:
            return func(sample)
    return wrapper
        

# FUNCTIONS {{{1
def isnumber(value):
    return float('-inf') < value < float('inf')


def isiterable(x):
    try:
        iter(x)
        return True
    except TypeError:
        return False


def fwhmToSigma(fwhm):
    """Calculate the standard deviation (sigma) from the Full width at Half Maximum of a gaussian"""
    return fwhm / (2 * math.sqrt(2 * math.log(2)))


def sigmatoFwhm(sigma):
    """Calculate the Full width at Half Maximum of a gaussian from the standard deviation (sigma)"""
    return sigma * 2 * math.sqrt(2 * math.log(2))


def calculateFullWidthAtNthMaximum(n, sigma):
    """Calculate the full width at nth maximum of a gaussian

    Usage

    fwnm = calculateFullWidthAtNthMaximum(n, sigma)

    Example
    -------
    # The full width at half maximum
    fwhm = calculateFullWidthAtNthMaximum(2, sigma)

    # The full width at tenth maximum
    fwtm = calculateFullWidthAtNthMaximum(10, sigma)

    """
    return 2 * math.sqrt(2 * math.log(n)) * sigma


@arenumbers
def differential(y):
    return map(lambda b, a: b - a, y[1:], y[0:-1])


def derivative(x, y):
    dy = differential(y)
    dx = differential(x)
    return avg(map(lambda a, b: a / b, dy, dx))


def linspace(start, stop, num=50, endpoint=True, retstep=False):
    if endpoint:
        step = float(stop - start) / (num - 1)
    else:
        step = float(stop - start) / num

    samples = [start + step * i for i in range(num)]

    if retstep:
        return samples, step
    else:
        return samples


@arenumbers
def squared(sample):
    return [x**2 for x in sample]


@arenumbers
def avg(sample):
    return float(sum(sample)) / len(sample)


@arenumbers
def std(sample):
    return math.sqrt(avg(squared(sample)) - avg(sample)**2)


@arenumbers
def _min(sample):
    return min(sample)


@arenumbers
def _max(sample):
    return max(sample)


@arenumbers
@elementwise
def _log(sample): 
    return math.log(sample)


@arenumbers
@elementwise
def _log10(sample): 
    return math.log10(sample)


@arenumbers
@elementwise
def _exp(sample): 
    return math.exp(sample)


@arenumbers
@elementwise
def _sin(sample): 
    return math.sin(sample)


@arenumbers
@elementwise
def _cos(sample): 
    return math.cos(sample)


@arenumbers
@elementwise
def _tan(sample): 
    return math.tan(sample)


@arenumbers
def _sum(sample):
    return sum(sample)


@arenumbers
def _cumsum(sample):
    return [sample[i] + _sum(sample[0:i]) for i in range(len(sample))]


# The root mean square
def rms(sample):
    return math.sqrt(avg(squared(sample)))


def snr(signal, noise):
    return rms(signal)**2 / rms(noise)**2


def snr_dB(signal, noise):
    return 10 * math.log10(snr(signal, noise))
    
    
def zipWith(f, a, b):
    """Zip a and b and apply f on each pair (a_i, b_i)
    
    Usage:
        c = zipWith(f, a, b)

    Examples:
        >>> a = [1, 10, 20]
        >>> b = [0, 9, 3]
        >>> zipWith(operator.add, a, b)
        [1, 19, 23]

        >>> zipWith(lambda x, y: x**2 + 3*y, a, b)
        [1, 127, 409]

    """
    return map(lambda x: f(x[0], x[1]), zip(a, b))


# def linear(x, xd, yd):
#     """Linear interpolation: return the ordinates at abscissa x 
    
#     Usage
#     -----
#     y = linear(x, xd, yd)

#     Parameters
#     ----------
#     x : iterable
#         The abscissa where y must be interpolated

#     xd : iterable
#         The abscissa of the data

#     yd : iterable
#         The abscissa of the data

#     Returns
#     -------
#     y : iterable
#         The interpolated ordinates
    
#     """
#     j = bisect.bisect_left(xd, x)
#     i = j - 1
#     if i < 0:
#         return yd[0]
#     if j >= len(xd):
#         return yd[-1]
        
#     return yd[i] + (yd[j] - yd[i]) / (xd[j] - xd[i]) * (x - xd[i])



# CLASSES {{{1

# TODO Improve the display of the string representation
# TODO Fix __setslice__

class Array(object):
    """A simple Array class roughly emulating nympy.array with array.
    
    Usage
    -----

    # Create an Array from an iterable
    >>> myarray = Array(iterable)

    # Create an empty Array from a number of values
    >>> myarray = Array(n)

    """
    MAX_DISPLAY_SIZE = 100
    DTYPE_CONVERSIONS = {int: "i", float: "d", bool: "B"}
    TYPE_CONVERSIONS  = dict((v, k) for k, v in DTYPE_CONVERSIONS.items())

    def __init__(self, values, dtype=None):
        if not dtype:
            self._dtype = self.__analyze_dtype(values)
        else:
            self._dtype = dtype

        self.values = array.array(self.dtype, self.__analyze_values(values))
        self.size = len(self.values)

    def get_dtype(self):
        return self._dtype

    def set_dtype(self, value):
        if value in (self.DTYPE_CONVERSIONS.values()):
            if value != self._dtype:
                self._dtype = value
                new_type = self.TYPE_CONVERSIONS[self._dtype]
                self.values = array.array(self._dtype, [new_type(x) for x in self.values])
        else:
            msg = "Valid dtypes are the values of: %r" % self.DTYPE_CONVERSIONS
            raise ValueError(msg)

    dtype = property(get_dtype, set_dtype)

    def __analyze_values(self, values):
        # An iterable is provided
        if isiterable(values):
            return values
        # A number of values is provided
        elif isinstance(values, int):
            return [0.0 for _ in range(values)]

    def __analyze_dtype(self, values):
        """The dtype is determined automatically from the first type in the values"""
        if isiterable(values):
            first_type = type(values[0])
            is_unique_type = all(x.__class__ == first_type for x in values)
            if not is_unique_type:
                msg = "All the values must have the same type."
                raise TypeError(msg)
            else:
                return self.DTYPE_CONVERSIONS[first_type]
        else:
            # Assume the user wants to initialize an Array of n floats
            # (where n = values) 
            return self.DTYPE_CONVERSIONS[float]
 
    def __str__(self):
        s = " ".join("%s" % str(x) for x in self.values)
        # if self.size <= self.MAX_DISPLAY_SIZE:
        #     s = " ".join("%s" % str(x) for x in self.values)
        # else:
        #     s = " ".join("%s" % str(x) for x in self.values if x <= self.values[2] or x >= self.values[-3])
        return "[%s]" % s

    __repr__ = __str__

    def __getitem__(self, i):
        # Return the ith+1 element
        if isinstance(i, int):
            return self.values[i]

        # Return the elements that satisfy the condition
        # Ex:
        # myarray[myarray > myarray.mean()]
        elif isinstance(i, self.__class__):
            indexes = [x for x in range(i.size) if i[x] == 1]
            return self.__class__(map(self.values.__getitem__, indexes))
        else:
            msg = "Use an index or a condition on the Array (ex: myarray[myarray > myarray.mean()])"
            raise IndexError(msg)

    def __setitem__(self, i, value):
        self.values[i] = value

    def __len__(self):
        return self.size

    def __perform_inverse_operation(self, operation, x):
        op = getattr(operator, operation)
        indexes = xrange(self.size)

        if isiterable(x):
            if len(x) == self.size:
                return self.__class__([op(x[i], self.values[i]) for i in indexes])
            else:
                msg = "The operands don't have the same size (%s) (%s)" % (self.size, len(x))
                raise ValueError(msg)
        else:
            return self.__class__([op(x, self.values[i]) for i in indexes])

    def __perform_operation(self, operation, x):
        op = getattr(operator, operation)
        indexes = xrange(self.size)

        if isiterable(x):
            if len(x) == self.size:
                return self.__class__([op(self.values[i], x[i]) for i in indexes])
            else:
                msg = "The operands don't have the same size (%d) (%d)" % (self.size, len(x))
                raise ValueError(msg)
        else:
            return self.__class__([op(self.values[i], x) for i in indexes])

    def __add__(self, x):
        return self.__perform_operation("add", x)

    def __sub__(self, x):
        return self.__perform_operation("sub", x)

    def __mul__(self, x):
        return self.__perform_operation("mul", x)

    def __rmul__(self, x):
        return self.__perform_inverse_operation("mul", x)

    def __div__(self, x):
        return self.__perform_operation("div", x)

    def __mod__(self, x):
        return self.__perform_operation("mod", x)

    def __rdiv__(self, x):
        return self.__perform_inverse_operation("div", x)

    def __eq__(self, x):
        return self.__perform_operation("eq", x)

    def __ne__(self, x):
        return self.__perform_operation("ne", x)

    def __lt__(self, x):
        return self.__perform_operation("lt", x)

    def __le__(self, x):
        return self.__perform_operation("le", x)

    def __gt__(self, x):
        return self.__perform_operation("gt", x)

    def __ge__(self, x):
        return self.__perform_operation("ge", x)

    def __neg__(self):
        return self.__perform_operation("mul", -1)

    def __pos__(self):
        return self

    def __or__(self, x):
        return self.__perform_operation("or", x)

    def __pow__(self, x):
        return self.__perform_operation("pow", x)

    def mean(self):
        return avg(self.values)

    def median(self):
        numbers = filter(isnumber, self.values)
        return median(numbers)

    def std(self):
        return std(self.values)

    def min(self):
        return _min(self.values)

    def max(self):
        return _max(self.values)

    def sum(self):
        return _sum(self.values)

    def cumsum(self):
        return Array(_cumsum(self.values))

    def copy(self):
        return copy.deepcopy(self)

    def plot(self, title="", xlabel="x", ylabel="y", legend=None, lineWidth=1, lineStyle="line", lineColor=None, plotType="line"):
        try:
            p = Plot(title)
            p.plot(range(self.size), self.values, 
                   xlabel=xlabel,
                   ylabel=ylabel,
                   legend=legend,
                   lineWidth=lineWidth,
                   lineStyle=lineStyle,
                   lineColor=lineColor,
                   plotType=plotType)
        except NameError:
            pass
        return p

    def count(self, value):
        """Return the number of occurrences of the given value"""
        if isnumber(value):
            return len([x for x in self.values if x == value])
        else:
            return len([x for x in self.values if not isnumber(x)])

    def __getslice__(self, imin, imax):
        return self.__class__(operator.getslice(self.values, imin, imax))

    def __setslice__(self, imin, imax, values):
        if isiterable(values):
            n = len(values)
            slice_size = imax-imin
            if n != slice_size:
                msg = "The input values (%d) don't have the same size as the slice (%d)" % (n, imax-imin)
                raise ValueError(msg)

            for i in range(imin, imax):
                self.values[i] = values[i]
        else:
            for i in range(imin, imax):
                self.values[i] = values

            return self.values


# vim: set fdm=marker:
