from __future__ import division, print_function, absolute_import
import pytest
import numpy as np
from numpy.testing import (TestCase, assert_array_almost_equal,
                           assert_array_equal, assert_, assert_allclose,
                           assert_equal)
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import LinearOperator
from scipy.optimize._differentiable_functions import (ScalarFunction,
                                                      VectorFunction,
                                                      LinearVectorFunction,
                                                      IdentityVectorFunction)
from scipy.optimize._hessian_update_strategy import BFGS


class ExScalarFunction:

    def __init__(self):
        self.nfev = 0
        self.ngev = 0
        self.nhev = 0

    def fun(self, x):
        self.nfev += 1
        return 2*(x[0]**2 + x[1]**2 - 1) - x[0]

    def grad(self, x):
        self.ngev += 1
        return np.array([4*x[0]-1, 4*x[1]])

    def hess(self, x):
        self.nhev += 1
        return 4*np.eye(2)


class TestScalarFunction(TestCase):

    def test_finite_difference_grad(self):
        ex = ExScalarFunction()
        nfev = 0
        ngev = 0

        x0 = [1.0, 0.0]
        analit = ScalarFunction(ex.fun, x0, (), ex.grad,
                                ex.hess, None, (-np.inf, np.inf))
        nfev += 1
        ngev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev, nfev)
        approx = ScalarFunction(ex.fun, x0, (), '2-point',
                                ex.hess, None, (-np.inf, np.inf))
        nfev += 3
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(analit.f, approx.f)
        assert_array_almost_equal(analit.g, approx.g)

        x = [10, 0.3]
        f_analit = analit.fun(x)
        g_analit = analit.grad(x)
        nfev += 1
        ngev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        f_approx = approx.fun(x)
        g_approx = approx.grad(x)
        nfev += 3
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_almost_equal(f_analit, f_approx)
        assert_array_almost_equal(g_analit, g_approx)

        x = [2.0, 1.0]
        g_analit = analit.grad(x)
        ngev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)

        g_approx = approx.grad(x)
        nfev += 3
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_almost_equal(g_analit, g_approx)

        x = [2.5, 0.3]
        f_analit = analit.fun(x)
        g_analit = analit.grad(x)
        nfev += 1
        ngev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        f_approx = approx.fun(x)
        g_approx = approx.grad(x)
        nfev += 3
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_almost_equal(f_analit, f_approx)
        assert_array_almost_equal(g_analit, g_approx)

        x = [2, 0.3]
        f_analit = analit.fun(x)
        g_analit = analit.grad(x)
        nfev += 1
        ngev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        f_approx = approx.fun(x)
        g_approx = approx.grad(x)
        nfev += 3
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_almost_equal(f_analit, f_approx)
        assert_array_almost_equal(g_analit, g_approx)

    def test_finite_difference_hess_linear_operator(self):
        ex = ExScalarFunction()
        nfev = 0
        ngev = 0
        nhev = 0

        x0 = [1.0, 0.0]
        analit = ScalarFunction(ex.fun, x0, (), ex.grad,
                                ex.hess, None, (-np.inf, np.inf))
        nfev += 1
        ngev += 1
        nhev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev, nhev)
        approx = ScalarFunction(ex.fun, x0, (), ex.grad,
                                '2-point', None, (-np.inf, np.inf))
        assert_(isinstance(approx.H, LinearOperator))
        for v in ([1.0, 2.0], [3.0, 4.0], [5.0, 2.0]):
            assert_array_equal(analit.f, approx.f)
            assert_array_almost_equal(analit.g, approx.g)
            assert_array_almost_equal(analit.H.dot(v), approx.H.dot(v))
        nfev += 1
        ngev += 4
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)

        x = [2.0, 1.0]
        H_analit = analit.hess(x)
        nhev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)
        H_approx = approx.hess(x)
        assert_(isinstance(H_approx, LinearOperator))
        for v in ([1.0, 2.0], [3.0, 4.0], [5.0, 2.0]):
            assert_array_almost_equal(H_analit.dot(v), H_approx.dot(v))
        ngev += 4
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)

        x = [2.1, 1.2]
        H_analit = analit.hess(x)
        nhev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)
        H_approx = approx.hess(x)
        assert_(isinstance(H_approx, LinearOperator))
        for v in ([1.0, 2.0], [3.0, 4.0], [5.0, 2.0]):
            assert_array_almost_equal(H_analit.dot(v), H_approx.dot(v))
        ngev += 4
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)

        x = [2.5, 0.3]
        _ = analit.grad(x)
        H_analit = analit.hess(x)
        ngev += 1
        nhev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)
        _ = approx.grad(x)
        H_approx = approx.hess(x)
        assert_(isinstance(H_approx, LinearOperator))
        for v in ([1.0, 2.0], [3.0, 4.0], [5.0, 2.0]):
            assert_array_almost_equal(H_analit.dot(v), H_approx.dot(v))
        ngev += 4
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)

        x = [5.2, 2.3]
        _ = analit.grad(x)
        H_analit = analit.hess(x)
        ngev += 1
        nhev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)
        _ = approx.grad(x)
        H_approx = approx.hess(x)
        assert_(isinstance(H_approx, LinearOperator))
        for v in ([1.0, 2.0], [3.0, 4.0], [5.0, 2.0]):
            assert_array_almost_equal(H_analit.dot(v), H_approx.dot(v))
        ngev += 4
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev+approx.nfev, nfev)
        assert_array_equal(ex.ngev, ngev)
        assert_array_equal(analit.ngev+approx.ngev, ngev)
        assert_array_equal(ex.nhev, nhev)
        assert_array_equal(analit.nhev+approx.nhev, nhev)

    def test_x_storage_overlap(self):
        # Scalar_Function should not store references to arrays, it should
        # store copies - this checks that updating an array in-place causes
        # Scalar_Function.x to be updated.

        def f(x):
            return np.sum(np.asarray(x) ** 2)

        x = np.array([1., 2., 3.])
        sf = ScalarFunction(f, x, (), '3-point', lambda x: x, None, (-np.inf, np.inf))

        assert x is not sf.x
        assert_equal(sf.fun(x), 14.0)
        assert x is not sf.x

        x[0] = 0.
        f1 = sf.fun(x)
        assert_equal(f1, 13.0)

        x[0] = 1
        f2 = sf.fun(x)
        assert_equal(f2, 14.0)
        assert x is not sf.x

        # now test with a HessianUpdate strategy specified
        hess = BFGS()
        x = np.array([1., 2., 3.])
        sf = ScalarFunction(f, x, (), '3-point', hess, None, (-np.inf, np.inf))

        assert x is not sf.x
        assert_equal(sf.fun(x), 14.0)
        assert x is not sf.x

        x[0] = 0.
        f1 = sf.fun(x)
        assert_equal(f1, 13.0)

        x[0] = 1
        f2 = sf.fun(x)
        assert_equal(f2, 14.0)
        assert x is not sf.x


class ExVectorialFunction:

    def __init__(self):
        self.nfev = 0
        self.njev = 0
        self.nhev = 0

    def fun(self, x):
        self.nfev += 1
        return np.array([2*(x[0]**2 + x[1]**2 - 1) - x[0],
                         4*(x[0]**3 + x[1]**2 - 4) - 3*x[0]])

    def jac(self, x):
        self.njev += 1
        return np.array([[4*x[0]-1, 4*x[1]],
                         [12*x[0]**2-3, 8*x[1]]])

    def hess(self, x, v):
        self.nhev += 1
        return v[0]*4*np.eye(2) + v[1]*np.array([[24*x[0], 0],
                                                 [0, 8]])


class TestVectorialFunction(TestCase):

    def test_finite_difference_jac(self):
        ex = ExVectorialFunction()
        nfev = 0
        njev = 0

        x0 = [1.0, 0.0]
        v0 = [0.0, 1.0]
        analit = VectorFunction(ex.fun, x0, ex.jac, ex.hess, None, None,
                                (-np.inf, np.inf), None)
        nfev += 1
        njev += 1
        assert_array_equal(ex.nfev, nfev)
        assert_array_equal(analit.nfev, nfev)
        assert_array_equal(ex.njev, njev)
        assert_array_equal(analit.njev, njev)
        approx = VectorFunction(ex.fun, x0, '2-point', ex.hess, None, None,
                                (-np.inf, np.inf), None)
        nfev += 3
        assert_array_equal(ex.nfev, nfev)