from __future__ import division, print_function, absolute_import

import numpy as np
from numpy.testing import assert_array_equal, assert_almost_equal, \
                          assert_array_almost_equal, assert_equal

from scipy.misc import pade, logsumexp, face, ascent


def test_pade_trivial():
    nump, denomp = pade([1.0], 0)
    assert_array_equal(nump.c, [1.0])
    assert_array_equal(denomp.c, [1.0])


def test_pade_4term_exp():
    # First four Taylor coefficients of exp(x).
    # Unlike poly1d, the first array element is the zero-order term.
    an = [1.0, 1.0, 0.5, 1.0/6]

    nump, denomp = pade(an, 0)
    assert_array_almost_equal(nump.c, [1.0/6, 0.5, 1.0, 1.0])
    assert_array_almost_equal(denomp.c, [1.0])

    nump, denomp = pade(an, 1)
    assert_array_almost_equal(nump.c, [1.0/6, 2.0/3, 1.0])
    assert_array_almost_equal(denomp.c, [-1.0/3, 1.0])

    nump, denomp = pade(an, 2)
    assert_array_almost_equal(nump.c, [1.0/3, 1.0])
    assert_array_almost_equal(denomp.c, [1.0/6, -2.0/3, 1.0])

    nump, denomp = pade(an, 3)
    assert_array_almost_equal(nump.c, [1.0])
    assert_array_almost_equal(denomp.c, [-1.0/6, 0.5, -1.0, 1.0])


def test_logsumexp():
    """Test whether logsumexp() function correctly handles large inputs."""
    a = np.arange(200)
    desired = np.log(np.sum(np.exp(a)))
    assert_almost_equal(logsumexp(a), desired)

    # Now test with large numbers
    b = [1000, 1000]
    desired = 1000.0 + np.log(2.0)
    assert_almost_equal(logsumexp(b), desired)

    n = 1000
    b = np.ones(n) * 10000
    desired = 10000.0 + np.log(n)
    assert_almost_equal(logsumexp(b), desired)

    x = np.array([1e-40] * 1000000)
    logx = np.log(x)

    X = np.vstack([x, x])
    logX = np.vstack([logx, logx])
    assert_array_almost_equal(np.exp(logsumexp(logX)), X.sum())
    assert_array_almost_equal(np.exp(logsumexp(logX, axis=0)), X.sum(axis=0))
    assert_array_almost_equal(np.exp(logsumexp(logX, axis=1)), X.sum(axis=1))


def test_logsumexp_b():
    a = np.arange(200)
    b = np.arange(200, 0, -1)
    desired = np.log(np.sum(b*np.exp(a)))
    assert_almost_equal(logsumexp(a, b=b), desired)

    a = [1000, 1000]
    b = [1.2, 1.2]
    desired = 1000 + np.log(2 * 1.2)
    assert_almost_equal(logsumexp(a, b=b), desired)

    x = np.array([1e-40] * 100000)
    b = np.linspace(1, 1000, 1e5)
    logx = np.log(x)

    X = np.vstack((x, x))
    logX = np.vstack((logx, logx))
    B = np.vstack((b, b))
    assert_array_almost_equal(np.exp(logsumexp(logX, b=B)), (B * X).sum())
    assert_array_almost_equal(np.exp(logsumexp(logX, b=B, axis=0)),
                                (B * X).sum(axis=0))
    assert_array_almost_equal(np.exp(logsumexp(logX, b=B, axis=1)),
                                (B * X).sum(axis=1))


def test_face():
    assert_equal(face().shape, (768, 1024, 3))


def test_ascent():
    assert_equal(ascent().shape, (512, 512))

@@ -268,16 +268,61 @@
 ascent%0A%0A
+from scipy.lib._version import NumpyVersion%0A%0A
 %0Adef tes
@@ -2008,32 +2008,1058 @@
 X.sum(axis=1))%0A%0A
+    # Handling special values properly%0A    assert_equal(logsumexp(np.inf), np.inf)%0A    assert_equal(logsumexp(-np.inf), -np.inf)%0A    assert_equal(logsumexp(np.nan), np.nan)%0A%0A    # Handling an array with different magnitudes on the axes%0A    assert_array_almost_equal(logsumexp(%5B%5B1e10, 1e-10%5D,%0A                                         %5B-1e10, -np.inf%5D%5D, axis=-1),%0A                              %5B1e10, -1e10%5D)%0A%0A    # Test keeping dimensions%0A    assert_array_almost_equal(logsumexp(%5B%5B1e10, 1e-10%5D,%0A                                         %5B-1e10, -np.inf%5D%5D, %0A                                        axis=-1,%0A                                        keepdims=True),%0A                              %5B%5B1e10%5D, %5B-1e10%5D%5D)%0A%0A    # Test multiple axes%0A    if NumpyVersion(np.version.version) %3E= NumpyVersion('1.7.0'):%0A        assert_array_almost_equal(logsumexp(%5B%5B1e10, 1e-10%5D,%0A                                             %5B-1e10, -np.inf%5D%5D, %0A                                            axis=(-1,-2)),%0A                                  1e10)%0A%0A
 %0Adef test_logsum

Add a few unit tests for the bug fixes in scipy.misc.logsumexp