'''
Bspines and smoothing splines.

General references:

    Craven, P. and Wahba, G. (1978) "Smoothing noisy data with spline functions.
    Estimating the correct degree of smoothing by
    the method of generalized cross-validation."
    Numerische Mathematik, 31(4), 377-403.

    Hastie, Tibshirani and Friedman (2001). "The Elements of Statistical
    Learning." Springer-Verlag. 536 pages.

    Hutchison, M. and Hoog, F. "Smoothing noisy data with spline functions."
    Numerische Mathematik, 47(1), 99-106.
'''

import numpy as np
import numpy.linalg as L

from scipy.linalg import solveh_banded
from scipy.optimize import golden
from models import _hbspline     #removed because this was segfaulting

# Issue warning regarding heavy development status of this module
import warnings
_msg = """
The bspline code is technology preview and requires significant work
on the public API and documentation. The API will likely change in the future
"""
warnings.warn(_msg, FutureWarning)


def _band2array(a, lower=0, symmetric=False, hermitian=False):
    """
    Take an upper or lower triangular banded matrix and return a
    numpy array.

    INPUTS:
       a         -- a matrix in upper or lower triangular banded matrix
       lower     -- is the matrix upper or lower triangular?
       symmetric -- if True, return the original result plus its transpose
       hermitian -- if True (and symmetric False), return the original
                    result plus its conjugate transposed
    """

    n = a.shape[1]
    r = a.shape[0]
    _a = 0

    if not lower:
        for j in range(r):
            _b = np.diag(a[r-1-j],k=j)[j:(n+j),j:(n+j)]
            _a += _b
            if symmetric and j > 0:
                _a += _b.T
            elif hermitian and j > 0:
                _a += _b.conjugate().T
    else:
        for j in range(r):
            _b = np.diag(a[j],k=j)[0:n,0:n]
            _a += _b
            if symmetric and j > 0:
                _a += _b.T
            elif hermitian and j > 0:
                _a += _b.conjugate().T
        _a = _a.T

    return _a


def _upper2lower(ub):
    """
    Convert upper triangular banded matrix to lower banded form.

    INPUTS:
       ub  -- an upper triangular banded matrix

    OUTPUTS: lb
       lb  -- a lower triangular banded matrix with same entries
              as ub
    """

    lb = np.zeros(ub.shape, ub.dtype)
    nrow, ncol = ub.shape
    for i in range(ub.shape[0]):
        lb[i,0:(ncol-i)] = ub[nrow-1-i,i:ncol]
        lb[i,(ncol-i):] = ub[nrow-1-i,0:i]
    return lb

def _lower2upper(lb):
    """
    Convert lower triangular banded matrix to upper banded form.

    INPUTS:
       lb  -- a lower triangular banded matrix

    OUTPUTS: ub
       ub  -- an upper triangular banded matrix with same entries
              as lb
    """

    ub = np.zeros(lb.shape, lb.dtype)
    nrow, ncol = lb.shape
    for i in range(lb.shape[0]):
        ub[nrow-1-i,i:ncol] = lb[i,0:(ncol-i)]
        ub[nrow-1-i,0:i] = lb[i,(ncol-i):]
    return ub

def _triangle2unit(tb, lower=0):
    """
    Take a banded triangular matrix and return its diagonal and the
    unit matrix: the banded triangular matrix with 1's on the diagonal,
    i.e. each row is divided by the corresponding entry on the diagonal.

    INPUTS:
       tb    -- a lower triangular banded matrix
       lower -- if True, then tb is assumed to be lower triangular banded,
                in which case return value is also lower triangular banded.

    OUTPUTS: d, b
       d     -- diagonal entries of tb
       b     -- unit matrix: if lower is False, b is upper triangular
                banded and its rows of have been divided by d,
                else lower is True, b is lower triangular banded
                and its columns have been divieed by d.
    """

    if lower:
        d = tb[0].copy()
    else:
        d = tb[-1].copy()

    if lower:
        return d, (tb / d)
    else:
        lnum = _upper2lower(tb)
        return d, _lower2upper(lnum / d)


def _trace_symbanded(a, b, lower=0):
    """
    Compute the trace(ab) for two upper or banded real symmetric matrices
    stored either in either upper or lower form.

    INPUTS:
       a, b    -- two banded real symmetric matrices (either lower or upper)
       lower   -- if True, a and b are assumed to be the lower half


    OUTPUTS: trace
       trace   -- trace(ab)
    """

    if lower:
        t = _zero_triband(a * b, lower=1)
        return t[0].sum() + 2 * t[1:].sum()
    else:
        t = _zero_triband(a * b, lower=0)
        return t[-1].sum() + 2 * t[:-1].sum()


def _zero_triband(a, lower=0):
    """
    Explicitly zero out unused elements of a real symmetric banded matrix.

    INPUTS:
       a   -- a real symmetric banded matrix (either upper or lower hald)
       lower   -- if True, a is assumed to be the lower half
    """

    nrow, ncol = a.shape
    if lower:
        for i in range(nrow):
            a[i, (ncol-i):] = 0.
    else:
        for i in range(nrow):
            a[i, 0:i] = 0.
    return a


class BSpline(object):

    '''

    Bsplines of a given order and specified knots.

    Implementation is based on description in Chapter 5 of

    Hastie, Tibshirani and Friedman (2001). "The Elements of Statistical
    Learning." Springer-Verlag. 536 pages.


    INPUTS:
       knots  -- a sorted array of knots with knots[0] the lower boundary,
                 knots[1] the upper boundary and knots[1:-1] the internal
                 knots.
       order  -- order of the Bspline, default is 4 which yields cubic
                 splines
       M      -- number of additional boundary knots, if None it defaults
                 to order
       coef   -- an optional array of real-valued coefficients for the Bspline
                 of shape (knots.shape + 2 * (M - 1) - order,).
       x      -- an optional set of x values at which to evaluate the
                 Bspline to avoid extra evaluation in the __call__ method

    '''
    # FIXME: update parameter names, replace single character names
    # FIXME: `order` should be actual spline order (implemented as order+1)
    ## FIXME: update the use of spline order in extension code (evaluate is recursively called)
    # FIXME: eliminate duplicate M and m attributes (m is order, M is related to tau size)

    def __init__(self, knots, order=4, M=None, coef=None, x=None):

        knots = np.squeeze(np.unique(np.asarray(knots)))

        if knots.ndim != 1:
            raise ValueError('expecting 1d array for knots')

        self.m = order
        if M is None:
            M = self.m
        self.M = M

        self.tau = np.hstack([[knots[0]]*(self.M-1), knots, [knots[-1]]*(self.M-1)])

        self.K = knots.shape[0] - 2
        if coef is None:
            self.coef = np.zeros((self.K + 2 * self.M - self.m), np.float64)
        else:
            self.coef = np.squeeze(coef)
            if self.coef.shape != (self.K + 2 * self.M - self.m):
                raise ValueError('coefficients of Bspline have incorrect shape')
        if x is not None:
            self.x = x

    def _setx(self, x):
        self._x = x
        self._basisx = self.basis(self._x)

    def _getx(self):
        return self._x

    x = property(_getx, _setx)

    def __call__(self, *args):
        """
        Evaluate the BSpline at a given point, yielding
        a matrix B and return

        B * self.coef


        INPUTS:
           args -- optional arguments. If None, it returns self._basisx,
                   the BSpline evaluated at the x values passed in __init__.
                   Otherwise, return the BSpline evaluated at the
                   first argument args[0].

        OUTPUTS: y
           y    -- value of Bspline at specified x values

        BUGS:
           If self has no attribute x, an exception will be raised
           because self has no attribute _basisx.
        """

        if not args:
            b = self._basisx.T
        else:
            x = args[0]
            b = np.asarray(self.basis(x)).T
        return np.squeeze(np.dot(b, self.coef))

    def basis_element(self, x, i, d=0):
        """
        Evaluate a particular basis element of the BSpline,
        or its derivative.

        INPUTS:
           x  -- x values at which to evaluate the basis element
           i  -- which element of the BSpline to return
           d  -- the order of derivative

        OUTPUTS: y
           y  -- value of d-th derivative of the i-th basis element
                 of the BSpline at specified x values
        """

        x = np.asarray(x, np.float64)
        _shape = x.shape
        if _shape == ():
            x.shape = (1,)
        x.shape = (np.product(_shape,axis=0),)
        if i < self.tau.shape[0] - 1:
            # TODO: OWNDATA flags...
            v = _hbspline.evaluate(x, self.tau, self.m, d, i, i+1)
        else:
            return np.zeros(x.shape, np.float64)

        if (i == self.tau.shape[0] - self.m):
            v = np.where(np.equal(x, self.tau[-1]), 1, v)
        v.shape = _shape
        return v

    def basis(self, x, d=0, lower=None, upper=None):
        """
        Evaluate the basis of the BSpline or its derivative.
        If lower or upper is specified, then only
        the [lower:upper] elements of the basis are returned.

        INPUTS:
           x     -- x values at which to evaluate the basis element
           i     -- which element of the BSpline to return
           d     -- the order of derivative
           lower -- optional lower limit of the set of basis
                    elements
           upper -- optional upper limit of the set of basis
                    elements

        OUTPUTS: y
           y  -- value of d-th derivative of the basis elements
                 of the BSpline at specified x values
        """
        x = np.asarray(x)
        _shape = x.shape
        if _shape == ():
            x.shape = (1,)
        x.shape = (np.product(_shape,axis=0),)

        if upper is None:
            upper = self.tau.shape[0] - self.m
        if lower is None:
            lower = 0
        upper = min(upper, self.tau.shape[0] - self.m)
        lower = max(0, lower)

        d = np.asarray(d)
        if d.shape == ():
            v = _hbspline.evaluate(x, self.tau, self.m, int(d), lower, upper)
        else:
            if d.shape[0] != 2:
                raise ValueError("if d is not an integer, expecting a jx2 \
                   array with first row indicating order \
                   of derivative, second row coefficient in front.")