from __future__ import division, print_function, absolute_import

import itertools

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

from numpy import mgrid, pi, sin, ogrid, poly1d, linspace
import numpy as np

from scipy._lib.six import xrange
from scipy._lib._numpy_compat import _assert_warns, suppress_warnings

from scipy.interpolate import (interp1d, interp2d, lagrange, PPoly, BPoly,
         splrep, splev, splantider, splint, sproot, Akima1DInterpolator,
         RegularGridInterpolator, LinearNDInterpolator, NearestNDInterpolator,
         RectBivariateSpline, interpn, NdPPoly, BSpline)

from scipy.special import poch, gamma

from scipy.interpolate import _ppoly

from scipy._lib._gcutils import assert_deallocated, IS_PYPY

from scipy.integrate import nquad

from scipy.special import binom

from scipy.sparse.sputils import matrix


class TestInterp2D(object):
    def test_interp2d(self):
        y, x = mgrid[0:2:20j, 0:pi:21j]
        z = sin(x+0.5*y)
        I = interp2d(x, y, z)
        assert_almost_equal(I(1.0, 2.0), sin(2.0), decimal=2)

        v,u = ogrid[0:2:24j, 0:pi:25j]
        assert_almost_equal(I(u.ravel(), v.ravel()), sin(u+0.5*v), decimal=2)

    def test_interp2d_meshgrid_input(self):
        # Ticket #703
        x = linspace(0, 2, 16)
        y = linspace(0, pi, 21)
        z = sin(x[None,:] + y[:,None]/2.)
        I = interp2d(x, y, z)
        assert_almost_equal(I(1.0, 2.0), sin(2.0), decimal=2)

    def test_interp2d_meshgrid_input_unsorted(self):
        np.random.seed(1234)
        x = linspace(0, 2, 16)
        y = linspace(0, pi, 21)

        z = sin(x[None,:] + y[:,None]/2.)
        ip1 = interp2d(x.copy(), y.copy(), z, kind='cubic')

        np.random.shuffle(x)
        z = sin(x[None,:] + y[:,None]/2.)
        ip2 = interp2d(x.copy(), y.copy(), z, kind='cubic')

        np.random.shuffle(x)
        np.random.shuffle(y)
        z = sin(x[None,:] + y[:,None]/2.)
        ip3 = interp2d(x, y, z, kind='cubic')

        x = linspace(0, 2, 31)
        y = linspace(0, pi, 30)

        assert_equal(ip1(x, y), ip2(x, y))
        assert_equal(ip1(x, y), ip3(x, y))

    def test_interp2d_eval_unsorted(self):
        y, x = mgrid[0:2:20j, 0:pi:21j]
        z = sin(x + 0.5*y)
        func = interp2d(x, y, z)

        xe = np.array([3, 4, 5])
        ye = np.array([5.3, 7.1])
        assert_allclose(func(xe, ye), func(xe, ye[::-1]))

        assert_raises(ValueError, func, xe, ye[::-1], 0, 0, True)

    def test_interp2d_linear(self):
        # Ticket #898
        a = np.zeros([5, 5])
        a[2, 2] = 1.0
        x = y = np.arange(5)
        b = interp2d(x, y, a, 'linear')
        assert_almost_equal(b(2.0, 1.5), np.array([0.5]), decimal=2)
        assert_almost_equal(b(2.0, 2.5), np.array([0.5]), decimal=2)

    def test_interp2d_bounds(self):
        x = np.linspace(0, 1, 5)
        y = np.linspace(0, 2, 7)
        z = x[None, :]**2 + y[:, None]

        ix = np.linspace(-1, 3, 31)
        iy = np.linspace(-1, 3, 33)

        b = interp2d(x, y, z, bounds_error=True)
        assert_raises(ValueError, b, ix, iy)

        b = interp2d(x, y, z, fill_value=np.nan)
        iz = b(ix, iy)
        mx = (ix < 0) | (ix > 1)
        my = (iy < 0) | (iy > 2)
        assert_(np.isnan(iz[my,:]).all())
        assert_(np.isnan(iz[:,mx]).all())
        assert_(np.isfinite(iz[~my,:][:,~mx]).all())


class TestInterp1D(object):

    def setup_method(self):
        self.x5 = np.arange(5.)
        self.x10 = np.arange(10.)
        self.y10 = np.arange(10.)
        self.x25 = self.x10.reshape((2,5))
        self.x2 = np.arange(2.)
        self.y2 = np.arange(2.)
        self.x1 = np.array([0.])
        self.y1 = np.array([0.])

        self.y210 = np.arange(20.).reshape((2, 10))
        self.y102 = np.arange(20.).reshape((10, 2))
        self.y225 = np.arange(20.).reshape((2, 2, 5))
        self.y25 = np.arange(10.).reshape((2, 5))
        self.y235 = np.arange(30.).reshape((2, 3, 5))
        self.y325 = np.arange(30.).reshape((3, 2, 5))

        self.fill_value = -100.0

    def test_validation(self):
        # Make sure that appropriate exceptions are raised when invalid values
        # are given to the constructor.

        # These should all work.
        for kind in ('nearest', 'zero', 'linear', 'slinear', 'quadratic',
                     'cubic', 'previous', 'next'):
            interp1d(self.x10, self.y10, kind=kind)
            interp1d(self.x10, self.y10, kind=kind, fill_value="extrapolate")
        interp1d(self.x10, self.y10, kind='linear', fill_value=(-1, 1))
        interp1d(self.x10, self.y10, kind='linear',
                 fill_value=np.array([-1]))
        interp1d(self.x10, self.y10, kind='linear',
                 fill_value=(-1,))
        interp1d(self.x10, self.y10, kind='linear',
                 fill_value=-1)
        interp1d(self.x10, self.y10, kind='linear',
                 fill_value=(-1, -1))
        interp1d(self.x10, self.y10, kind=0)
        interp1d(self.x10, self.y10, kind=1)
        interp1d(self.x10, self.y10, kind=2)
        interp1d(self.x10, self.y10, kind=3)
        interp1d(self.x10, self.y210, kind='linear', axis=-1,
                 fill_value=(-1, -1))
        interp1d(self.x2, self.y210, kind='linear', axis=0,
                 fill_value=np.ones(10))
        interp1d(self.x2, self.y210, kind='linear', axis=0,
                 fill_value=(np.ones(10), np.ones(10)))
        interp1d(self.x2, self.y210, kind='linear', axis=0,
                 fill_value=(np.ones(10), -1))

        # x array must be 1D.
        assert_raises(ValueError, interp1d, self.x25, self.y10)

        # y array cannot be a scalar.
        assert_raises(ValueError, interp1d, self.x10, np.array(0))

        # Check for x and y arrays having the same length.
        assert_raises(ValueError, interp1d, self.x10, self.y2)
        assert_raises(ValueError, interp1d, self.x2, self.y10)
        assert_raises(ValueError, interp1d, self.x10, self.y102)
        interp1d(self.x10, self.y210)
        interp1d(self.x10, self.y102, axis=0)

        # Check for x and y having at least 1 element.
        assert_raises(ValueError, interp1d, self.x1, self.y10)
        assert_raises(ValueError, interp1d, self.x10, self.y1)
        assert_raises(ValueError, interp1d, self.x1, self.y1)

        # Bad fill values
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=(-1, -1, -1))  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=[-1, -1, -1])  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=np.array((-1, -1, -1)))  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=[[-1]])  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=[-1, -1])  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=np.array([]))  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x10, self.y10, kind='linear',
                      fill_value=())  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x2, self.y210, kind='linear',
                      axis=0, fill_value=[-1, -1])  # doesn't broadcast
        assert_raises(ValueError, interp1d, self.x2, self.y210, kind='linear',
                      axis=0, fill_value=(0., [-1, -1]))  # above doesn't bc

    def test_init(self):
        # Check that the attributes are initialized appropriately by the
        # constructor.
        assert_(interp1d(self.x10, self.y10).copy)
        assert_(not interp1d(self.x10, self.y10, copy=False).copy)
        assert_(interp1d(self.x10, self.y10).bounds_error)
        assert_(not interp1d(self.x10, self.y10, bounds_error=False).bounds_error)
        assert_(np.isnan(interp1d(self.x10, self.y10).fill_value))
        assert_equal(interp1d(self.x10, self.y10, fill_value=3.0).fill_value,
                     3.0)
        assert_equal(interp1d(self.x10, self.y10, fill_value=(1.0, 2.0)).fill_value,
                     (1.0, 2.0))
        assert_equal(interp1d(self.x10, self.y10).axis, 0)
        assert_equal(interp1d(self.x10, self.y210).axis, 1)
        assert_equal(interp1d(self.x10, self.y102, axis=0).axis, 0)
        assert_array_equal(interp1d(self.x10, self.y10).x, self.x10)
        assert_array_equal(interp1d(self.x10, self.y10).y, self.y10)
        assert_array_equal(interp1d(self.x10, self.y210).y, self.y210)

    def test_assume_sorted(self):
        # Check for unsorted arrays
        interp10 = interp1d(self.x10, self.y10)
        interp10_unsorted = interp1d(self.x10[::-1], self.y10[::-1])

        assert_array_almost_equal(interp10_unsorted(self.x10), self.y10)
        assert_array_almost_equal(interp10_unsorted(1.2), np.array([1.2]))
        assert_array_almost_equal(interp10_unsorted([2.4, 5.6, 6.0]),
                                  interp10([2.4, 5.6, 6.0]))

        # Check assume_sorted keyword (defaults to False)
        interp10_assume_kw = interp1d(self.x10[::-1], self.y10[::-1],
                                      assume_sorted=False)
        assert_array_almost_equal(interp10_assume_kw(self.x10), self.y10)

        interp10_assume_kw2 = interp1d(self.x10[::-1], self.y10[::-1],
                                       assume_sorted=True)
        # Should raise an error for unsorted input if assume_sorted=True
        assert_raises(ValueError, interp10_assume_kw2, self.x10)

        # Check that if y is a 2-D array, things are still consistent
        interp10_y_2d = interp1d(self.x10, self.y210)
        interp10_y_2d_unsorted = interp1d(self.x10[::-1], self.y210[:, ::-1])
        assert_array_almost_equal(interp10_y_2d(self.x10),
                                  interp10_y_2d_unsorted(self.x10))

    def test_linear(self):
        for kind in ['linear', 'slinear']:
            self._check_linear(kind)

    def _check_linear(self, kind):
        # Check the actual implementation of linear interpolation.
        interp10 = interp1d(self.x10, self.y10, kind=kind)
        assert_array_almost_equal(interp10(self.x10), self.y10)
        assert_array_almost_equal(interp10(1.2), np.array([1.2]))
        assert_array_almost_equal(interp10([2.4, 5.6, 6.0]),
                                  np.array([2.4, 5.6, 6.0]))

        # test fill_value="extrapolate"
        extrapolator = interp1d(self.x10, self.y10, kind=kind,
                                fill_value='extrapolate')
        assert_allclose(extrapolator([-1., 0, 9, 11]),
                        [-1, 0, 9, 11], rtol=1e-14)

        opts = dict(kind=kind,
                    fill_value='extrapolate',
                    bounds_error=True)
        assert_raises(ValueError, interp1d, self.x10, self.y10, **opts)

    def test_linear_dtypes(self):
        # regression test for gh-5898, where 1D linear interpolation has been
        # delegated to numpy.interp for all float dtypes, and the latter was
        # not handling e.g. np.float128.
        for dtyp in np.sctypes["float"]:
            x = np.arange(8, dtype=dtyp)
            y = x
            yp = interp1d(x, y, kind='linear')(x)
            assert_equal(yp.dtype, dtyp)
            assert_allclose(yp, y, atol=1e-15)

    def test_slinear_dtypes(self):
        # regression test for gh-7273: 1D slinear interpolation fails with
        # float32 inputs
        dt_r = [np.float16, np.float32, np.float64]
        dt_rc = dt_r + [np.complex64, np.complex128]
        spline_kinds = ['slinear', 'zero', 'quadratic', 'cubic']
        for dtx in dt_r:
            x = np.arange(0, 10, dtype=dtx)
            for dty in dt_rc:
                y = np.exp(-x/3.0).astype(dty)
                for dtn in dt_r:
                    xnew = x.astype(dtn)
                    for kind in spline_kinds:
                        f = interp1d(x, y, kind=kind, bounds_error=False)
                        assert_allclose(f(xnew), y, atol=1e-7,
                                        err_msg="%s, %s %s" % (dtx, dty, dtn))

    def test_cubic(self):
        # Check the actual implementation of spline interpolation.
        interp10 = interp1d(self.x10, self.y10, kind='cubic')
        assert_array_almost_equal(interp10(self.x10), self.y10)
        assert_array_almost_equal(interp10(1.2), np.array([1.2]))
        assert_array_almost_equal(interp10([2.4, 5.6, 6.0]),
                                  np.array([2.4, 5.6, 6.0]),)

    def test_nearest(self):
        # Check the actual implementation of nearest-neighbour interpolation.
        interp10 = interp1d(self.x10, self.y10, kind='nearest')
        assert_array_almost_equal(interp10(self.x10), self.y10)
        assert_array_almost_equal(interp10(1.2), np.array(1.))
        assert_array_almost_equal(interp10([2.4, 5.6, 6.0]),
                                  np.array([2., 6., 6.]),)

        # test fill_value="extrapolate"
        extrapolator = interp1d(self.x10, self.y10, kind='nearest',
                                fill_value='extrapolate')
        assert_allclose(extrapolator([-1., 0, 9, 11]),
                        [0, 0, 9, 9], rtol=1e-14)

        opts = dict(kind='nearest',
                    fill_value='extrapolate',
                    bounds_error=True)
        assert_raises(ValueError, interp1d, self.x10, self.y10, **opts)

    def test_previous(self):
        # Check the actual implementation of previous interpolation.
        interp10 = interp1d(self.x10, self.y10, kind='previous')
        assert_array_almost_equal(interp10(self.x10), self.y10)
        assert_array_almost_equal(interp10(1.2), np.array(1.))
        assert_array_almost_equal(interp10([2.4, 5.6, 6.0]),
                                  np.array([2., 5., 6.]),)

        # test fill_value="extrapolate"
        extrapolator = interp1d(self.x10, self.y10, kind='previous',
                                fill_value='extrapolate')
        assert_allclose(extrapolator([-1., 0, 9, 11]),
                        [0, 0, 9, 9], rtol=1e-14)

        opts = dict(kind='previous',
                    fill_value='extrapolate',
                    bounds_error=True)