# Source code for numdifftools.tests.test_nd_algopy

    # -*- coding:utf-8 -*-
    """"""

    from __future__ import division
    import unittest
    import numdifftools.nd_algopy as nd
    import numpy as np
    from numpy import pi, r_, sqrt, array
    from numpy.testing import assert_array_almost_equal
    from scipy import linalg, optimize, constants

    _TINY = np.finfo(float).machar.tiny

    #  Hamiltonian
    #     H = sum_i(p_i2/(2m)+ 1/2 * m * w2 x_i2) + sum_(i!=j)(a/|x_i-x_j|)
    class classicalHamiltonian(object):
        """
        Hamiltonian

        Parameters
        ----------
        N : scalar
            number of ions in the chain
        w : scalar
            angular trap frequency
        C : scalar
            Coulomb constant times the electronic charge in SI units.
        m : scalar
            the mass of a single trapped ion in the chain
        """

        def __init__(self):
            self.N = 2
            f = 1000000      # f is a scalar, it's the trap frequency
            self.w = 2 * pi * f
            self.C = (4 * pi * constants.epsilon_0) ** (-1) * constants.e ** 2
            # C is a scalar, it's the I
            self.m = 39.96 * 1.66e-27

        def potential(self, positionvector):
            """
            Return potential

            positionvector is an 1-d array (vector) of length N that contains the
            positions of the N ions
            """
            x = positionvector
            w = self.w
            C = self.C
            m = self.m

            # First we consider the potential of the harmonic oscillator
            Vx = 0.5 * m * (w ** 2) * sum(x ** 2)
            # then we add the coulomb interaction:
            for i, xi in enumerate(x):
                for xj in x[i + 1:]:
                    Vx += C / (abs(xi - xj))
            return Vx

        def initialposition(self):
            """Defines initial position as an estimate for the minimize process."""
            N = self.N
            x_0 = r_[-(N - 1) / 2:(N - 1) / 2:N * 1j]
            return x_0

        def normal_modes(self, eigenvalues):
            """Return normal modes

            Computed eigenvalues of the matrix Vx are of the form
                (normal_modes)**2*m.
            """
            m = self.m
            normal_modes = sqrt(eigenvalues / m)
            return normal_modes

    def _run_hamiltonian(verbose=True):
        c = classicalHamiltonian()
        if verbose:
            print(c.potential(array([-0.5, 0.5])))
            print(c.potential(array([-0.5, 0.0])))
            print(c.potential(array([0.0, 0.0])))

        xopt = optimize.fmin(c.potential, c.initialposition(), xtol=1e-10)

        hessian = nd.Hessian(c.potential)

        H = hessian(xopt)
        true_H = np.array([[5.23748385e-12, -2.61873829e-12],
                           [-2.61873829e-12, 5.23748385e-12]])
        error_estimate = np.NAN
        if verbose:
            print(xopt)
            print('H', H)
            print('H-true_H', np.abs(H - true_H))
            # print('error_estimate', info.error_estimate)

            eigenvalues = linalg.eigvals(H)
            normal_modes = c.normal_modes(eigenvalues)

            print('eigenvalues', eigenvalues)
            print('normal_modes', normal_modes)
        return H, error_estimate, true_H

    class TestHessian(unittest.TestCase):

        def test_run_hamiltonian(self):
            H, _error_estimate, true_H = _run_hamiltonian(verbose=False)
            self.assertTrue((np.abs(H - true_H) < 1e-18).all())

        def test_hessian_cosIx_yI_at_I0_0I(self):
            # cos(x-y), at (0,0)

            def fun(xy):
                return np.cos(xy[0] - xy[1])
            htrue = [[-1., 1.], [1., -1.]]
            methods = ['forward', ]  # 'reverse']

            for method in methods:
                Hfun2 = nd.Hessian(fun, method=method)
                h2 = Hfun2([0, 0])
                # print(method, (h2-np.array(htrue)))
                assert_array_almost_equal(h2, htrue)

    class TestDerivative(unittest.TestCase):

        # TODO: Derivative does not tackle non-finite values.
        #     def test_infinite_functions(self):
        #         def finf(x):
        #             return np.inf * np.ones_like(x)
        #         df = nd.Derivative(finf, method='forward')
        #         val = df(0)
        #         self.assert_(np.isnan(val))

        def test_high_order_derivative_cos(self):
            true_vals = (-1.0, 0.0, 1.0, 0.0) * 5

            x = np.pi / 2  # np.linspace(0, np.pi/2, 15)
            for method in ['forward', 'reverse']:
                nmax = 15 if method in ['forward'] else 2
                for n in range(1, nmax):
                    d3cos = nd.Derivative(np.cos, n=n, method=method)
                    y = d3cos(x)
                    assert_array_almost_equal(y, true_vals[n - 1])

        def test_fun_with_additional_parameters(self):
            """Test for issue #9"""
            def func(x, a, b=1):
                return b * a * x * x * x
            methods = ['reverse', 'forward']
            dfuns = [nd.Jacobian, nd.Derivative, nd.Gradient,  nd.Hessdiag,
                     nd.Hessian]
            for dfun in dfuns:
                for method in methods:
                    df = dfun(func, method=method)
                    val = df(0.0, 1.0, b=2)
                    assert_array_almost_equal(val, 0)

        def test_derivative_cube(self):
            """Test for Issue 7"""
            def cube(x):
                return x * x * x

            shape = (3, 2)
            x = np.ones(shape) * 2
            for method in ['forward', 'reverse']:
                dcube = nd.Derivative(cube, method=method)
                dx = dcube(x)
                assert_array_almost_equal(list(dx.shape), list(shape),
                                          decimal=13,
                                          err_msg='Shape mismatch')
                txt = 'First differing element %d\n value = %g,\n true value = %g'
                for i, (val, tval) in enumerate(zip(dx.ravel(),
                                                    (3 * x**2).ravel())):
                    assert_array_almost_equal(val, tval, decimal=8,
                                              err_msg=txt % (i, val, tval))

        def test_derivative_exp(self):
            # derivative of exp(x), at x == 0
            for method in ['forward', 'reverse']:
                dexp = nd.Derivative(np.exp, method=method)
                assert_array_almost_equal(dexp(0), np.exp(0), decimal=8)

        def test_derivative_sin(self):
            # Evaluate the indicated (default = first)
            # derivative at multiple points
            for method in ['forward', 'reverse']:
                dsin = nd.Derivative(np.sin, method=method)
                x = np.linspace(0, 2. * np.pi, 13)
                y = dsin(x)
                np.testing.assert_almost_equal(y, np.cos(x), decimal=8)

        def test_derivative_on_sinh(self):
            for method in ['forward', ]:  # 'reverse']: # TODO: reverse fails
                dsinh = nd.Derivative(np.sinh, method=method)
                self.assertAlmostEqual(dsinh(0.0), np.cosh(0.0))

        def test_derivative_on_log(self):

            x = np.r_[0.01, 0.1]
            for method in ['forward', 'reverse']:
                dlog = nd.Derivative(np.log, method=method)

                assert_array_almost_equal(dlog(x), 1.0 / x)

    class TestJacobian(unittest.TestCase):

        def test_on_scalar_function(self):
            def f2(x):
                return x[0] * x[1] * x[2] + np.exp(x[0]) * x[1]
            for method in ['forward', 'reverse']:
                Jfun3 = nd.Jacobian(f2, method=method)
                x = Jfun3([3., 5., 7.])
                assert_array_almost_equal(x, [[135.42768462, 41.08553692, 15.]])

        def test_on_vector_valued_function(self):
            xdata = np.reshape(np.arange(0, 1, 0.1), (-1, 1))
            ydata = 1 + 2 * np.exp(0.75 * xdata)

            def fun(c):
                return (c[0] + c[1] * np.exp(c[2] * xdata) - ydata) ** 2

            for method in ['reverse']:  # TODO: 'forward' fails

                Jfun = nd.Jacobian(fun, method=method)
                J = Jfun([1, 2, 0.75])  # should be numerically zero
                assert_array_almost_equal(J, np.zeros(J.shape))

    class TestGradient(unittest.TestCase):

        def test_on_scalar_function(self):
            def fun(x):
                return np.sum(x ** 2)

            dtrue = [2., 4., 6.]

            for method in ['forward', 'reverse']:  #

                dfun = nd.Gradient(fun, method=method)
                d = dfun([1, 2, 3])
                assert_array_almost_equal(d, dtrue)

    class TestHessdiag(unittest.TestCase):

        def test_forward(self):
            def fun(x):
                return x[0] + x[1] ** 2 + x[2] ** 3
            htrue = np.array([0., 2., 18.])
            Hfun = nd.Hessdiag(fun)
            hd = Hfun([1, 2, 3])
            _error = hd - htrue
            assert_array_almost_equal(hd, htrue)

        def test_reverse(self):
            def fun(x):
                return x[0] + x[1] ** 2 + x[2] ** 3
            htrue = np.array([0., 2., 18.])
            Hfun = nd.Hessdiag(fun, method='reverse')
            hd = Hfun([1, 2, 3])
            _error = hd - htrue
            assert_array_almost_equal(hd, htrue)

    if __name__ == '__main__':
        # _run_hamiltonian()
        unittest.main()