from __future__ import division, print_function, absolute_import

import warnings

import numpy as np

from numpy.testing import (
    assert_almost_equal, assert_array_equal, assert_array_almost_equal,
    assert_allclose, assert_equal, assert_)
from pytest import raises as assert_raises

from scipy.interpolate import (
    KroghInterpolator, krogh_interpolate,
    BarycentricInterpolator, barycentric_interpolate,
    approximate_taylor_polynomial, CubicHermiteSpline, pchip,
    PchipInterpolator, pchip_interpolate, Akima1DInterpolator, CubicSpline,
    make_interp_spline)

from scipy._lib.six import xrange


def check_shape(interpolator_cls, x_shape, y_shape, deriv_shape=None, axis=0,
                extra_args={}):
    np.random.seed(1234)

    x = [-1, 0, 1, 2, 3, 4]
    s = list(range(1, len(y_shape)+1))
    s.insert(axis % (len(y_shape)+1), 0)
    y = np.random.rand(*((6,) + y_shape)).transpose(s)

    # Cython code chokes on y.shape = (0, 3) etc, skip them
    if y.size == 0:
        return

    xi = np.zeros(x_shape)
    if interpolator_cls is CubicHermiteSpline:
        dydx = np.random.rand(*((6,) + y_shape)).transpose(s)
        yi = interpolator_cls(x, y, dydx, axis=axis, **extra_args)(xi)
    else:
        yi = interpolator_cls(x, y, axis=axis, **extra_args)(xi)

    target_shape = ((deriv_shape or ()) + y.shape[:axis]
                    + x_shape + y.shape[axis:][1:])
    assert_equal(yi.shape, target_shape)

    # check it works also with lists
    if x_shape and y.size > 0:
        if interpolator_cls is CubicHermiteSpline:
            interpolator_cls(list(x), list(y), list(dydx), axis=axis,
                             **extra_args)(list(xi))
        else:
            interpolator_cls(list(x), list(y), axis=axis,
                             **extra_args)(list(xi))

    # check also values
    if xi.size > 0 and deriv_shape is None:
        bs_shape = y.shape[:axis] + (1,)*len(x_shape) + y.shape[axis:][1:]
        yv = y[((slice(None,),)*(axis % y.ndim)) + (1,)]
        yv = yv.reshape(bs_shape)

        yi, y = np.broadcast_arrays(yi, yv)
        assert_allclose(yi, y)


SHAPES = [(), (0,), (1,), (6, 2, 5)]


def test_shapes():

    def spl_interp(x, y, axis):
        return make_interp_spline(x, y, axis=axis)

    for ip in [KroghInterpolator, BarycentricInterpolator, CubicHermiteSpline,
               pchip, Akima1DInterpolator, CubicSpline, spl_interp]:
        for s1 in SHAPES:
            for s2 in SHAPES:
                for axis in range(-len(s2), len(s2)):
                    if ip != CubicSpline:
                        check_shape(ip, s1, s2, None, axis)
                    else:
                        for bc in ['natural', 'clamped']:
                            extra = {'bc_type': bc}
                            check_shape(ip, s1, s2, None, axis, extra)

def test_derivs_shapes():
    def krogh_derivs(x, y, axis=0):
        return KroghInterpolator(x, y, axis).derivatives

    for s1 in SHAPES:
        for s2 in SHAPES:
            for axis in range(-len(s2), len(s2)):
                check_shape(krogh_derivs, s1, s2, (6,), axis)


def test_deriv_shapes():
    def krogh_deriv(x, y, axis=0):
        return KroghInterpolator(x, y, axis).derivative

    def pchip_deriv(x, y, axis=0):
        return pchip(x, y, axis).derivative()

    def pchip_deriv2(x, y, axis=0):
        return pchip(x, y, axis).derivative(2)

    def pchip_antideriv(x, y, axis=0):
        return pchip(x, y, axis).derivative()

    def pchip_antideriv2(x, y, axis=0):
        return pchip(x, y, axis).derivative(2)

    def pchip_deriv_inplace(x, y, axis=0):
        class P(PchipInterpolator):
            def __call__(self, x):
                return PchipInterpolator.__call__(self, x, 1)
            pass
        return P(x, y, axis)

    def akima_deriv(x, y, axis=0):
        return Akima1DInterpolator(x, y, axis).derivative()

    def akima_antideriv(x, y, axis=0):
        return Akima1DInterpolator(x, y, axis).antiderivative()

    def cspline_deriv(x, y, axis=0):
        return CubicSpline(x, y, axis).derivative()

    def cspline_antideriv(x, y, axis=0):
        return CubicSpline(x, y, axis).antiderivative()

    def bspl_deriv(x, y, axis=0):
        return make_interp_spline(x, y, axis=axis).derivative()

    def bspl_antideriv(x, y, axis=0):
        return make_interp_spline(x, y, axis=axis).antiderivative()

    for ip in [krogh_deriv, pchip_deriv, pchip_deriv2, pchip_deriv_inplace,
               pchip_antideriv, pchip_antideriv2, akima_deriv, akima_antideriv,
               cspline_deriv, cspline_antideriv, bspl_deriv, bspl_antideriv]:
        for s1 in SHAPES:
            for s2 in SHAPES:
                for axis in range(-len(s2), len(s2)):
                    check_shape(ip, s1, s2, (), axis)


def test_complex():
    x = [1, 2, 3, 4]
    y = [1, 2, 1j, 3]

    for ip in [KroghInterpolator, BarycentricInterpolator, pchip, CubicSpline]:
        p = ip(x, y)
        assert_allclose(y, p(x))

    dydx = [0, -1j, 2, 3j]
    p = CubicHermiteSpline(x, y, dydx)
    assert_allclose(y, p(x))
    assert_allclose(dydx, p(x, 1))


class TestKrogh(object):
    def setup_method(self):
        self.true_poly = np.poly1d([-2,3,1,5,-4])
        self.test_xs = np.linspace(-1,1,100)
        self.xs = np.linspace(-1,1,5)
        self.ys = self.true_poly(self.xs)

    def test_lagrange(self):
        P = KroghInterpolator(self.xs,self.ys)
        assert_almost_equal(self.true_poly(self.test_xs),P(self.test_xs))

    def test_scalar(self):
        P = KroghInterpolator(self.xs,self.ys)
        assert_almost_equal(self.true_poly(7),P(7))
        assert_almost_equal(self.true_poly(np.array(7)), P(np.array(7)))

    def test_derivatives(self):
        P = KroghInterpolator(self.xs,self.ys)
        D = P.derivatives(self.test_xs)
        for i in xrange(D.shape[0]):
            assert_almost_equal(self.true_poly.deriv(i)(self.test_xs),
                                D[i])

    def test_low_derivatives(self):
        P = KroghInterpolator(self.xs,self.ys)
        D = P.derivatives(self.test_xs,len(self.xs)+2)
        for i in xrange(D.shape[0]):
            assert_almost_equal(self.true_poly.deriv(i)(self.test_xs),
                                D[i])

    def test_derivative(self):
        P = KroghInterpolator(self.xs,self.ys)
        m = 10
        r = P.derivatives(self.test_xs,m)
        for i in xrange(m):
            assert_almost_equal(P.derivative(self.test_xs,i),r[i])

    def test_high_derivative(self):
        P = KroghInterpolator(self.xs,self.ys)
        for i in xrange(len(self.xs),2*len(self.xs)):
            assert_almost_equal(P.derivative(self.test_xs,i),
                                np.zeros(len(self.test_xs)))

    def test_hermite(self):
        xs = [0,0,0,1,1,1,2]
        ys = [self.true_poly(0),
              self.true_poly.deriv(1)(0),
              self.true_poly.deriv(2)(0),
              self.true_poly(1),
              self.true_poly.deriv(1)(1),
              self.true_poly.deriv(2)(1),
              self.true_poly(2)]
        P = KroghInterpolator(self.xs,self.ys)
        assert_almost_equal(self.true_poly(self.test_xs),P(self.test_xs))

    def test_vector(self):
        xs = [0, 1, 2]
        ys = np.array([[0,1],[1,0],[2,1]])
        P = KroghInterpolator(xs,ys)
        Pi = [KroghInterpolator(xs,ys[:,i]) for i in xrange(ys.shape[1])]
        test_xs = np.linspace(-1,3,100)
        assert_almost_equal(P(test_xs),
                np.rollaxis(np.asarray([p(test_xs) for p in Pi]),-1))
        assert_almost_equal(P.derivatives(test_xs),
                np.transpose(np.asarray([p.derivatives(test_xs) for p in Pi]),
                    (1,2,0)))

    def test_empty(self):
        P = KroghInterpolator(self.xs,self.ys)
        assert_array_equal(P([]), [])

    def test_shapes_scalarvalue(self):
        P = KroghInterpolator(self.xs,self.ys)
        assert_array_equal(np.shape(P(0)), ())
        assert_array_equal(np.shape(P(np.array(0))), ())
        assert_array_equal(np.shape(P([0])), (1,))
        assert_array_equal(np.shape(P([0,1])), (2,))

    def test_shapes_scalarvalue_derivative(self):
        P = KroghInterpolator(self.xs,self.ys)
        n = P.n
        assert_array_equal(np.shape(P.derivatives(0)), (n,))
        assert_array_equal(np.shape(P.derivatives(np.array(0))), (n,))
        assert_array_equal(np.shape(P.derivatives([0])), (n,1))
        assert_array_equal(np.shape(P.derivatives([0,1])), (n,2))

    def test_shapes_vectorvalue(self):
        P = KroghInterpolator(self.xs,np.outer(self.ys,np.arange(3)))
        assert_array_equal(np.shape(P(0)), (3,))
        assert_array_equal(np.shape(P([0])), (1,3))
        assert_array_equal(np.shape(P([0,1])), (2,3))

    def test_shapes_1d_vectorvalue(self):
        P = KroghInterpolator(self.xs,np.outer(self.ys,[1]))
        assert_array_equal(np.shape(P(0)), (1,))
        assert_array_equal(np.shape(P([0])), (1,1))
        assert_array_equal(np.shape(P([0,1])), (2,1))

    def test_shapes_vectorvalue_derivative(self):
        P = KroghInterpolator(self.xs,np.outer(self.ys,np.arange(3)))
        n = P.n
        assert_array_equal(np.shape(P.derivatives(0)), (n,3))
        assert_array_equal(np.shape(P.derivatives([0])), (n,1,3))
        assert_array_equal(np.shape(P.derivatives([0,1])), (n,2,3))

    def test_wrapper(self):
        P = KroghInterpolator(self.xs, self.ys)
        ki = krogh_interpolate
        assert_almost_equal(P(self.test_xs), ki(self.xs, self.ys, self.test_xs))
        assert_almost_equal(P.derivative(self.test_xs, 2),
                            ki(self.xs, self.ys, self.test_xs, der=2))
        assert_almost_equal(P.derivatives(self.test_xs, 2),
                            ki(self.xs, self.ys, self.test_xs, der=[0, 1]))

    def test_int_inputs(self):
        # Check input args are cast correctly to floats, gh-3669
        x = [0, 234, 468, 702, 936, 1170, 1404, 2340, 3744, 6084, 8424,
             13104, 60000]
        offset_cdf = np.array([-0.95, -0.86114777, -0.8147762, -0.64072425,
                               -0.48002351, -0.34925329, -0.26503107,
                               -0.13148093, -0.12988833, -0.12979296,
                               -0.12973574, -0.08582937, 0.05])
        f = KroghInterpolator(x, offset_cdf)

        assert_allclose(abs((f(x) - offset_cdf) / f.derivative(x, 1)),
                        0, atol=1e-10)

    def test_derivatives_complex(self):
        # regression test for gh-7381: krogh.derivatives(0) fails complex y
        x, y = np.array([-1, -1, 0, 1, 1]), np.array([1, 1.0j, 0, -1, 1.0j])
        func = KroghInterpolator(x, y)
        cmplx = func.derivatives(0)

        cmplx2 = (KroghInterpolator(x, y.real).derivatives(0) +
                  1j*KroghInterpolator(x, y.imag).derivatives(0))
        assert_allclose(cmplx, cmplx2, atol=1e-15)


class TestTaylor(object):
    def test_exponential(self):
        degree = 5
        p = approximate_taylor_polynomial(np.exp, 0, degree, 1, 15)
        for i in xrange(degree+1):
            assert_almost_equal(p(0),1)
            p = p.deriv()
        assert_almost_equal(p(0),0)


class TestBarycentric(object):
    def setup_method(self):
        self.true_poly = np.poly1d([-2, 3, 1, 5, -4])
        self.test_xs = np.linspace(-1, 1, 100)
        self.xs = np.linspace(-1, 1, 5)
        self.ys = self.true_poly(self.xs)

    def test_lagrange(self):
        P = BarycentricInterpolator(self.xs, self.ys)
        assert_almost_equal(self.true_poly(self.test_xs), P(self.test_xs))

    def test_scalar(self):
        P = BarycentricInterpolator(self.xs, self.ys)
        assert_almost_equal(self.true_poly(7), P(7))
        assert_almost_equal(self.true_poly(np.array(7)), P(np.array(7)))

    def test_delayed(self):
        P = BarycentricInterpolator(self.xs)
        P.set_yi(self.ys)
        assert_almost_equal(self.true_poly(self.test_xs), P(self.test_xs))

    def test_append(self):
        P = BarycentricInterpolator(self.xs[:3], self.ys[:3])
        P.add_xi(self.xs[3:], self.ys[3:])
        assert_almost_equal(self.true_poly(self.test_xs), P(self.test_xs))

    def test_vector(self):
        xs = [0, 1, 2]
        ys = np.array([[0, 1], [1, 0], [2, 1]])
        BI = BarycentricInterpolator 
        P = BI(xs, ys)
        Pi = [BI(xs, ys[:, i]) for i in xrange(ys.shape[1])]
        test_xs = np.linspace(-1, 3, 100)
        assert_almost_equal(P(test_xs),
                np.rollaxis(np.asarray([p(test_xs) for p in Pi]), -1))

    def test_shapes_scalarvalue(self):
        P = BarycentricInterpolator(self.xs, self.ys)
        assert_array_equal(np.shape(P(0)), ())