# -*- coding: utf-8 -*-
"""Python/Numpy implementation of Bspline basis functions via Cox - de Boor algorithm."""

from __future__ import division

from functools import partial
import numpy as np

class memoize(object):
    """
       Cache the return value of a method

       This class is meant to be used as a decorator of methods. The return value
       from a given method invocation will be cached on the instance whose method
       was invoked. All arguments passed to a method decorated with memoize must
       be hashable.

       If a memoized method is invoked directly on its class the result will not
       be cached. Instead the method will be invoked like a static method:
       class Obj(object):
           @memoize
           def add_to(self, arg):
               return self + arg
       Obj.add_to(1) # not enough arguments
       Obj.add_to(1, 2) # returns 3, result is not cached

       Script borrowed from here:
       MIT Licensed, attributed to Daniel Miller, Wed, 3 Nov 2010
       http://code.activestate.com/recipes/577452-a-memoize-decorator-for-instance-methods/
    """
    def __init__(self, func):
        self.func = func
    def __get__(self, obj, objtype=None):
        if obj is None:
            return self.func
        return partial(self, obj)
    def __call__(self, *args, **kw):
        obj = args[0]
        try:
            cache = obj.__cache
        except AttributeError:
            cache = obj.__cache = {}
        key = (self.func, args[1:], frozenset(kw.items()))
        try:
            res = cache[key]
        except KeyError:
            res = cache[key] = self.func(*args, **kw)
        return res



class Bspline():
    """
    Numpy implementation of Cox - de Boor algorithm in 1D

    inputs:
        knot_vector: Python list or Numpy array containing knot vector
                        entries
        order: Order of interpolation, e.g. 0 -> piecewise constant between
                  knots, 1 -> piecewise linear between knots, etc.
    last : bool
        if True the xi equal to the value of the last knot are also taken along

    outputs:
           basis object that is callable to evaluate basis functions at given
           values of knot span
    """

    def __init__(self, knot_vector, order, last=False ):
        """Initialize attributes"""
        self.knot_vector = np.array(knot_vector)
        if last :
            self.last_knot = np.full_like( knot_vector, knot_vector[0] - 1 )
            self.last_knot[-order-1] = knot_vector[-order-1]
        self.last = last
        self.p = order

        #Dummy calls to the functions for memory storage
        self.__call__(0.0)
        self.d(0.0)


    def __basis0(self, xi):
        """Order zero basis"""

        upper =  xi < self.knot_vector[1:]
        if self.last :
            upper = np.any( [upper, xi == self.last_knot[1:]], axis=0 )

        b = np.where(np.all([self.knot_vector[:-1] <= xi, upper], axis=0 ), 1.0, 0.0 )
        return b

    def __basis(self, xi, p, compute_derivatives=False):
        """
        Recursive Cox - de Boor function to compute basis functions and
           optionally their derivatives
        """

        if p == 0:
            return self.__basis0(xi)
        else:
            basis_p_minus_1 = self.__basis(xi, p - 1)

        first_term_numerator = xi - self.knot_vector[:-p]
        first_term_denominator = self.knot_vector[p:] - self.knot_vector[:-p]

        second_term_numerator = self.knot_vector[(p + 1):] - xi
        second_term_denominator = (self.knot_vector[(p + 1):] -
                                   self.knot_vector[1:-p])


        #Change numerator in last recursion if derivatives are desired
        if compute_derivatives and p == self.p:

            first_term_numerator = p
            second_term_numerator = -p

        #Disable divide by zero error because we check for it
        with np.errstate(divide='ignore', invalid='ignore'):
            first_term = np.where(first_term_denominator != 0.0,
                                  (first_term_numerator /
                                   first_term_denominator), 0.0)
            second_term = np.where(second_term_denominator != 0.0,
                                   (second_term_numerator /
                                    second_term_denominator), 0.0)

        return  (first_term[:-1] * basis_p_minus_1[:-1] +
                 second_term * basis_p_minus_1[1:])


    @memoize
    def __call__(self, xi):
        """
        Convenience function to make the object callable.  Also 'memoized'
           for speed.
        """
        return self.__basis(xi, self.p, compute_derivatives=False)

    @memoize
    def d(self, xi):
        """
        Convenience function to compute derivate of basis functions.
           'Memoized' for speed.
        """
        return self.__basis(xi, self.p, compute_derivatives=True)

    def plot(self):
        """
        Convenience function to plot basis functions over full
        range of knots.
        """

        import matplotlib.pyplot as plt

        x_min = np.min(self.knot_vector)
        x_max = np.max(self.knot_vector)

        x = np.linspace(x_min, x_max, num=1000)

        N = np.array([self(i) for i in x]).T;

        for n in N:

            plt.plot(x,n)

        return plt.show()

    def dplot(self):
        """
        Convenience function to plot derivatives of basis functions over
        full range of knots.
        """

        import matplotlib.pyplot as plt

        x_min = np.min(self.knot_vector)
        x_max = np.max(self.knot_vector)

        x = np.linspace(x_min, x_max, num=1000)

        N = np.array([self.d(i) for i in x]).T;

        for n in N:

            plt.plot(x,n)

        return plt.show()


    def __diff_internal(self):
        """
        Differentiate a B-spline once, and return the resulting coefficients and Bspline objects.

        This preserves the Bspline object nature of the data, enabling recursive implementation
        of higher-order differentiation (see `diff`).

        The value of the first derivative of `B` at a point `x` can be obtained as::

            def diff1(B, x):
                terms = B.__diff_internal()
                return sum( ci*Bi(x) for ci,Bi in terms )

        Returns:
            tuple of tuples, where each item is (coefficient, Bspline object).

        See:
            `diff`: differentiation of any order >= 0
        """
        assert self.p > 0, "order of Bspline must be > 0"  # we already handle the other case in diff()

        # https://www.cs.mtu.edu/~shene/COURSES/cs3621/NOTES/spline/B-spline/bspline-derv.html
        #
        t    = self.knot_vector
        p    = self.p
        Bi   = Bspline( t[:-1], p-1 )
        Bip1 = Bspline( t[1:],  p-1 )

        numer1 = +p
        numer2 = -p
        denom1 = t[p:-1]   - t[:-(p+1)]
        denom2 = t[(p+1):] - t[1:-p]

        with np.errstate(divide='ignore', invalid='ignore'):
            ci   = np.where(denom1 != 0., (numer1 / denom1), 0.)
            cip1 = np.where(denom2 != 0., (numer2 / denom2), 0.)

        return ( (ci,Bi), (cip1,Bip1) )


    def diff(self, order=1):
        """
        Differentiate a B-spline `order` number of times.

        Parameters:
        order: int, >= 0

        Returns:
            **lambda** `x`: ... that evaluates the `order`-th derivative of `B` at the point `x`.
        """
        order = int(order)
        if order < 0:
            raise ValueError("order must be >= 0, got %d" % (order))

        if order == 0:
            return self.__call__

        if order > self.p:   # identically zero, but force the same output format as in the general case
            dummy = self.__call__(0.)  # get number of basis functions and output dtype
            nbasis = dummy.shape[0]
            return lambda x: np.zeros( (nbasis,), dtype=dummy.dtype )  # accept but ignore input x

        # At each differentiation, each term maps into two new terms.
        # The number of terms in the result will be 2**order.
        #
        # This will cause an exponential explosion in the number of terms for high derivative orders,
        # but for the first few orders (practical usage; >3 is rarely needed) the approach works.
        #
        terms = [ (1.,self) ]
        for k in range(order):
            tmp = []
            for Ci,Bi in terms:
                tmp.extend( (Ci*cn, Bn) for cn,Bn in Bi.__diff_internal() )  # NOTE: also propagate Ci
            terms = tmp

        # perform final summation at call time
        return lambda x: sum( ci*Bi(x) for ci,Bi in terms )


    def collmat(self, tau, deriv_order=0):
        """
        Compute collocation matrix.

        Parameters:
        tau:
            Python list or rank-1 array, collocation sites
        deriv_order:
            int, >=0, order of derivative for which to compute the collocation matrix.
            The default is 0, which means the function value itself.

        Returns:
        A:
            rank-2 array such that
                A[i,j] = D**deriv_order B_j(tau[i])
            where
                D**k  = kth derivative (0 for function value itself)

        Example:
        If the coefficients of a spline function are given in the vector c, then::

            sum( A[i,:] * c[:] )

        will give a rank-1 array of function values at the sites tau[i] that were supplied
        to `collmat`.

        Similarly for derivatives (if the supplied `deriv_order`> 0).

        """
        # get number of basis functions and output dtype
        dummy = self.__call__(0.)
        nbasis = dummy.shape[0]

        tau = np.atleast_1d(tau)
        if tau.ndim > 1:
            raise ValueError("tau must be a list or a rank-1 array")

        A = np.empty( (tau.shape[0], nbasis), dtype=dummy.dtype )
        f = self.diff(order=deriv_order)
        for i,taui in enumerate(tau):
            A[i,:] = f(taui)

        return np.squeeze(A)
