# Source code for numdifftools.nd_algopy

    """
    Numdifftools.nd_algopy
    ======================
    This module provide an easy to use interface to derivatives calculated with
    AlgoPy. Algopy stands for Algorithmic Differentiation in Python.

    The purpose of AlgoPy is the evaluation of higher-order derivatives in the
    forward and reverse mode of Algorithmic Differentiation (AD) of functions that
    are implemented as Python programs. Particular focus are functions that contain
    numerical linear algebra functions as they often appear in statistically
    motivated functions. The intended use of AlgoPy is for easy prototyping at
    reasonable execution speeds. More precisely, for a typical program a
    directional derivative takes order 10 times as much time as time as the
    function evaluation. This is approximately also true for the gradient.

    Algoritmic differentiation
    ==========================

    Algorithmic differentiation (AD) is a set of techniques to numerically
    evaluate the derivative of a function specified by a computer program. AD
    exploits the fact that every computer program, no matter how complicated,
    executes a sequence of elementary arithmetic operations (addition,
    subtraction, multiplication, division, etc.) and elementary functions
    (exp, log, sin, cos, etc.). By applying the chain rule repeatedly to these
    operations, derivatives of arbitrary order can be computed automatically,
    accurately to working precision, and using at most a small constant factor
    more arithmetic operations than the original program.

    Algorithmic differentiation is not:

    Symbolic differentiation, nor Numerical differentiation (the method of
    finite differences). These classical methods run into problems:
    symbolic differentiation leads to inefficient code (unless carefully done)
    and faces the difficulty of converting a computer program into a single
    expression, while numerical differentiation can introduce round-off errors
    in the discretization process and cancellation. Both classical methods have
    problems with calculating higher derivatives, where the complexity and
    errors increase. Finally, both classical methods are slow at computing the
    partial derivatives of a function with respect to many inputs, as is needed
    for gradient-based optimization algorithms. Algoritmic differentiation
    solves all of these problems.

    Reference
    ---------
    Sebastian F. Walter and Lutz Lehmann 2013,
    "Algorithmic differentiation in Python with AlgoPy",
    in Journal of Computational Science, vol 4, no 5, pp 334 - 344,
    http://www.sciencedirect.com/science/article/pii/S1877750311001013

    https://en.wikipedia.org/wiki/Automatic_differentiation

    https://pythonhosted.org/algopy/index.html
    """
    from __future__ import division
    import numpy as np
    from scipy import misc
    try:
        import algopy
        from algopy import UTPM
    except ImportError:
        algopy = None

    _cmn_doc = """
        Calculate %(derivative)s with Algorithmic Differentiation method

        Parameters
        ----------
        f : function
            function of one array f(x, `*args`, `**kwds`)%(extra_parameter)s
        method : string, optional {'forward', 'reverse'}
            defines method used in the approximation
        %(returns)s
        Notes
        -----
        Algorithmic differentiation is a set of techniques to numerically
        evaluate the derivative of a function specified by a computer program. AD
        exploits the fact that every computer program, no matter how complicated,
        executes a sequence of elementary arithmetic operations (addition,
        subtraction, multiplication, division, etc.) and elementary functions
        (exp, log, sin, cos, etc.). By applying the chain rule repeatedly to these
        operations, derivatives of arbitrary order can be computed automatically,
        accurately to working precision, and using at most a small constant factor
        more arithmetic operations than the original program.
        %(extra_note)s
        Reference
        ---------
        Sebastian F. Walter and Lutz Lehmann 2013,
        "Algorithmic differentiation in Python with AlgoPy",
        in Journal of Computational Science, vol 4, no 5, pp 334 - 344,
        http://www.sciencedirect.com/science/article/pii/S1877750311001013

        https://en.wikipedia.org/wiki/Automatic_differentiation

        %(example)s
        %(see_also)s
        """

    class _Common(object):
        def __init__(self, f, method='forward'):
            self.f = f
            self.method = method
            self.n = 1
            self._computational_graph = None

        @property
        def f(self):
            return self._f

        @f.setter
        def f(self, f):
            self._f = f
            self._computational_graph = None

        def computational_graph(self, x, *args, **kwds):
            if self._computational_graph is None:
                # STEP 1: trace the function evaluation
                cg = algopy.CGraph()
                x = algopy.Function(x)

                y = self.f(x, *args, **kwds)
                # y = UTPM.as_utpm(z)
                cg.trace_off()
                cg.independentFunctionList = [x]
                cg.dependentFunctionList = [y]
                self._computational_graph = cg
            return self._computational_graph

        def _get_function(self):
            if self.n == 0:
                return self.f
            name = '_' + self.method
            return getattr(self, name)

        def __call__(self, x, *args, **kwds):
            fun = self._get_function()
            x0 = np.asarray(x, dtype=float)
            return fun(x0, *args, **kwds)

    class Derivative(_Common):
        __doc__ = _cmn_doc % dict(
            derivative='n-th derivative',
            extra_parameter="""
        n : int, optional
            Order of the derivative.""",
            extra_note="", returns="""
        Returns
        -------
        der : ndarray
           array of derivatives
        """, example="""
        Example
        -------
        # 1'st and 2'nd derivative of exp(x), at x == 1

        >>> import numpy as np
        >>> import numdifftools.nd_algopy as nda
        >>> fd = nda.Derivative(np.exp)              # 1'st derivative
        >>> np.allclose(fd(1), 2.718281828459045)
        True
        >>> fd5 = nda.Derivative(np.exp, n=5)         # 5'th derivative
        >>> np.allclose(fd5(1), 2.718281828459045)
        True

        # 1'st derivative of x^3+x^4, at x = [0,1]

        >>> f = lambda x: x**3 + x**4
        >>> fd3 = nda.Derivative(f)
        >>> np.allclose(fd3([0,1]), [ 0.,  7.])
        True
        """, see_also="""
        See also
        --------
        Gradient,
        Hessdiag,
        Hessian,
        Jacobian
        """)

        def __init__(self, f, n=1, method='forward'):
            self.f = f
            self.n = n
            self.method = method

        def _forward(self, x, *args, **kwds):
            x0 = np.asarray(x)
            shape = x0.shape
            P = 1
            x = UTPM(np.zeros((self.n + 1, P) + shape))
            x.data[0, 0] = x0
            x.data[1, 0] = 1
            z = self.f(x, *args, **kwds)
            y = UTPM.as_utpm(z)

            return y.data[self.n, 0] * misc.factorial(self.n)

        def _reverse(self, x, *args, **kwds):
            if self.n != 1:
                raise NotImplementedError('Derivative reverse not implemented'
                                          ' for n>1')

            c_graph = self.computational_graph(np.asarray(1), *args, **kwds)
            shape0 = x.shape
            y = np.array([c_graph.gradient(xi) for xi in x.ravel()])
            return y.reshape(shape0)

    class Jacobian(_Common):
        __doc__ = _cmn_doc % dict(
            derivative='Jacobian',
            extra_parameter="",
            extra_note="", returns="""
        Returns
        -------
        jacob : array
            Jacobian
        """, example="""
        Example
        -------
        >>> import numdifftools.nd_algopy as nda

        #(nonlinear least squares)

        >>> xdata = np.reshape(np.arange(0,1,0.1),(-1,1))
        >>> ydata = 1+2*np.exp(0.75*xdata)
        >>> f = lambda c: (c[0]+c[1]*np.exp(c[2]*xdata) - ydata)**2

        Jfun = nda.Jacobian(f) # Todo: This does not work
        Jfun([1,2,0.75]) # should be numerically zero
        array([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
               [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
               [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

        >>> Jfun2 = nda.Jacobian(f, method='reverse')
        >>> Jfun2([1,2,0.75]).T # should be numerically zero
        array([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
               [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
               [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

        >>> f2 = lambda x : x[0]*x[1]*x[2] + np.exp(x[0])*x[1]
        >>> Jfun3 = nda.Jacobian(f2)
        >>> Jfun3([3.,5.,7.])
        array([[ 135.42768462,   41.08553692,   15.        ]])

        >>> Jfun4 = nda.Jacobian(f2, method='reverse')
        >>> Jfun4([3,5,7])
        array([[ 135.42768462,   41.08553692,   15.        ]])
        """, see_also="""
        See also
        --------
        Derivative
        Gradient,
        Hessdiag,
        Hessian,
        """)

        #     def _jacobian_forward(self, x, *args, **kwds):
        #         x = np.asarray(x, dtype=float)
        #         # shape = x.shape
        #         D, Nm = 2, x.size
        #         P = Nm
        #         y = UTPM(np.zeros((D, P, Nm)))
        #
        #         y.data[0, :] = x.ravel()
        #         y.data[1, :] = np.eye(Nm)
        #         z0 = self.f(y, *args, **kwds)
        #         z = UTPM.as_utpm(z0)
        #         J = z.data[1, :, :, 0]
        #         return J

        def _forward(self, x, *args, **kwds):
            # forward mode without building the computational graph
            tmp = algopy.UTPM.init_jacobian(x)
            y = self.f(tmp, *args, **kwds)
            return np.atleast_2d(algopy.UTPM.extract_jacobian(y))

        def _reverse(self, x, *args, **kwds):
            x = np.atleast_1d(x)
            c_graph = self.computational_graph(x, *args, **kwds)
            return c_graph.jacobian(x)

    class Gradient(_Common):
        __doc__ = _cmn_doc % dict(
            derivative='Gradient',
            extra_parameter="",
            extra_note="", returns="""
        Returns
        -------
        grad : array
            gradient
        """, example="""
        Example
        -------
        >>> import numdifftools.nd_algopy as nda
        >>> f = lambda x: np.sum(x**2)
        >>> df = nda.Gradient(f, method='reverse')
        >>> df([1,2,3])
        array([ 2.,  4.,  6.])

        #At [x,y] = [1,1], compute the numerical gradient
        #of the function sin(x-y) + y*exp(x)

        >>> sin = np.sin; exp = np.exp
        >>> z = lambda xy: sin(xy[0]-xy[1]) + xy[1]*exp(xy[0])
        >>> dz = nda.Gradient(z)
        >>> grad2 = dz([1, 1])
        >>> grad2
        array([ 3.71828183,  1.71828183])

        #At the global minimizer (1,1) of the Rosenbrock function,
        #compute the gradient. It should be essentially zero.

        >>> rosen = lambda x : (1-x[0])**2 + 105.*(x[1]-x[0]**2)**2
        >>> rd = nda.Gradient(rosen)
        >>> grad3 = rd([1,1])
        >>> grad3==np.array([ 0.,  0.])
        array([ True,  True], dtype=bool)
        """, see_also="""
        See also
        --------
        Derivative
        Jacobian,
        Hessdiag,
        Hessian,
        """)

        def _reverse(self, x, *args, **kwds):

            c_graph = self.computational_graph(x, *args, **kwds)
            return c_graph.gradient(x)

        def _forward(self, x, *args, **kwds):
            # forward mode without building the computational graph

            tmp = algopy.UTPM.init_jacobian(x)
            y = self.f(tmp, *args, **kwds)
            return algopy.UTPM.extract_jacobian(y)

    class Hessian(_Common):
        __doc__ = _cmn_doc % dict(
            derivative='Hessian',
            extra_parameter="",
            returns="""
        Returns
        -------
        hess : ndarray
           array of partial second derivatives, Hessian
        """, extra_note='', example="""
        Example
        -------
        >>> import numdifftools.nd_algopy as nda

        # Rosenbrock function, minimized at [1,1]

        >>> rosen = lambda x : (1.-x[0])**2 + 105*(x[1]-x[0]**2)**2
        >>> Hf = nda.Hessian(rosen)
        >>> h = Hf([1, 1]) #  h =[ 842 -420; -420, 210];
        >>> h
        array([[ 842., -420.],
               [-420.,  210.]])

        # cos(x-y), at (0,0)

        >>> cos = np.cos
        >>> f = lambda xy : cos(xy[0]-xy[1])
        >>> Hfun2 = nda.Hessian(f)
        >>> h2 = Hfun2([0, 0]) # h2 = [-1 1; 1 -1]
        >>> h2
        array([[-1.,  1.],
               [ 1., -1.]])

        >>> Hfun3 = nda.Hessian(f, method='reverse')
        >>> h3 = Hfun3([0, 0]) # h2 = [-1, 1; 1, -1];
        >>> h3
        array([[-1.,  1.],
               [ 1., -1.]])
        """, see_also="""
        See also
        --------
        Derivative
        Gradient,
        Jacobian,
        Hessdiag,
        """)

        def __init__(self, f, method='forward'):
            self.f = f
            self.method = method
            self.n = 2
            self._computational_graph = None

        def _forward(self, x, *args, **kwds):
            x = np.atleast_1d(x)
            tmp = algopy.UTPM.init_hessian(x)
            y = self.f(tmp, *args, **kwds)
            return algopy.UTPM.extract_hessian(len(x), y)

        def _reverse(self, x, *args, **kwds):
            x = np.atleast_1d(np.asarray(x, dtype=float))
            c_graph = self.computational_graph(x, *args, **kwds)
            return c_graph.hessian(x)

    class Hessdiag(Hessian):
        __doc__ = _cmn_doc % dict(
            derivative='Hessian diagonal',
            extra_parameter="",
            returns="""
        Returns
        -------
        hessdiag : ndarray
           Hessian diagonal array of partial second order derivatives.
        """, extra_note='', example="""
        Example
        -------
        >>> import numdifftools.nd_algopy as nda

        # Rosenbrock function, minimized at [1,1]

        >>> rosen = lambda x : (1.-x[0])**2 + 105*(x[1]-x[0]**2)**2
        >>> Hfun = nda.Hessdiag(rosen)
        >>> h = Hfun([1, 1]) #  h =[ 842, 210]
        >>> h
        array([ 842.,  210.])

        # cos(x-y), at (0,0)

        >>> cos = np.cos
        >>> f = lambda xy : cos(xy[0]-xy[1])
        >>> Hfun2 = nda.Hessdiag(f)
        >>> h2 = Hfun2([0, 0]) # h2 = [-1, -1]
        >>> h2
        array([-1., -1.])

        >>> Hfun3 = nda.Hessdiag(f, method='reverse')
        >>> h3 = Hfun3([0, 0]) # h2 = [-1, -1];
        >>> h3
        array([-1., -1.])

        """, see_also="""
        See also
        --------
        Derivative
        Gradient,
        Jacobian,
        Hessian,
        """)

        def _forward(self, x, *args, **kwds):
            # return np.diag(super(Hessdiag, self)._forward(x, *args, **kwds))
            D, Nm = 2+1, x.size
            P = Nm
            y = UTPM(np.zeros((D, P, Nm)))

            y.data[0, :] = x.ravel()
            y.data[1, :] = np.eye(Nm)
            z0 = self.f(y, *args, **kwds)
            z = UTPM.as_utpm(z0)
            H = z.data[2, ...] * 2
            return H

        def _reverse(self, x, *args, **kwds):
            return np.diag(super(Hessdiag, self)._reverse(x, *args, **kwds))

    # def _example_taylor():
    #     def f(x):
    #         return x*x*x*x  # np.sin(np.cos(x) + np.sin(x))
    #     D = 5
    #     P = 1
    #     x = UTPM(np.zeros((D, P)))
    #     x.data[0, 0] = 1.0
    #     x.data[1, 0] = 1
    #
    #     y = f(x)
    #     print('coefficients of y =', y.data[:, 0])

    def test_docstrings():
        import doctest
        doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)

    if __name__ == '__main__':
        test_docstrings()
        # _example_taylor()