# Copyright 2014-2020 The ODL contributors
#
# This file is part of ODL.
#
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
# obtain one at https://mozilla.org/MPL/2.0/.

"""Unit tests for `diff_ops`."""

from __future__ import division

import pytest

import odl
from odl.core.discr.diff_ops import (
    Divergence, Gradient, Laplacian, PartialDerivative, finite_diff)
from odl.core.util.testutils import (
    all_almost_equal, all_equal, dtype_tol, noise_element, simple_fixture)
from odl.core.array_API_support import get_array_and_backend, odl_all_equal

# --- pytest fixtures --- #


method = simple_fixture('method', ['central', 'forward', 'backward'])
padding = simple_fixture('padding', [('constant', 0), ('constant', 1),
                                     'symmetric', 'periodic',
                                     'order0', 'order1', 'order2'])


@pytest.fixture(scope="module", params=[1, 2, 3], ids=['1d', '2d', '3d'])
def space(request, odl_impl_device_pairs):
    impl, device = odl_impl_device_pairs
    ndim = request.param

    return odl.uniform_discr([0] * ndim, [1] * ndim, [5] * ndim, impl=impl, device=device)

@pytest.fixture(scope="module")
def data(odl_impl_device_pairs):
    impl, device = odl_impl_device_pairs
    return odl.uniform_discr(0, 1, 10, impl=impl, device=device).element([0.5, 1, 3.5, 2, -.5, 3, -1, -1, 0, 3])
# --- finite_diff --- #


def test_finite_diff_invalid_args(data):
    """Test finite difference function for invalid arguments."""
    arr, backend = get_array_and_backend(data)
    # Test that old "edge order" argument fails.
    with pytest.raises(TypeError):
        finite_diff(data, axis=0, edge_order=0)

    # at least a two-element array is required
    with pytest.raises(ValueError):
        finite_diff(backend.array_constructor([0.0]), axis=0)

    # axis
    with pytest.raises(IndexError):
        finite_diff(data, axis=2)

    # in-place argument
    # size is an attribute of numpy arrays but of method of pytorch tensors
    out = backend.array_namespace.zeros(len(arr) + 1)
    with pytest.raises(ValueError):
        finite_diff(data, axis=0, out=out)
    with pytest.raises(ValueError):
        finite_diff(data, axis=0, dx=0)

    # wrong method
    with pytest.raises(ValueError):
        finite_diff(data, axis=0, method='non-method')


def test_finite_diff_explicit(data):
    """Compare finite differences function to explicit computation."""

    # phantom data
    arr, backend = get_array_and_backend(data)
    ns = backend.array_namespace

    # explicitly calculated finite difference
    diff_ex = ns.zeros_like(arr)

    # interior: second-order accurate differences
    diff_ex[1:-1] = (arr[2:] - arr[:-2]) / 2.0

    # default: out=None, axis=0, dx=1.0, zero_padding=None, method='forward'
    diff = finite_diff(arr, axis=0, dx=1.0, out=None,
                       pad_mode='constant')
    assert odl_all_equal(diff, finite_diff(arr, axis=0))

    # boundary: one-sided second-order accurate forward/backward difference
    diff = finite_diff(arr, axis=0, dx=1.0, out=None,
                       method='central', pad_mode='order2')
    diff_ex[0] = -(3 * arr[0] - 4 * arr[1] + arr[2]) / 2.0
    diff_ex[-1] = (3 * arr[-1] - 4 * arr[-2] + arr[-3]) / 2.0
    assert odl_all_equal(diff, diff_ex)

    # non-unit step length
    dx = 0.5
    diff = finite_diff(arr, axis=0, dx=dx, method='central', out=None,
                       pad_mode='order2')
    assert odl_all_equal(diff, diff_ex / dx)

    # boundary: second-order accurate central differences with zero padding
    diff = finite_diff(arr, axis=0, method='central', pad_mode='constant',
                       pad_const=0)
    diff_ex[0] = arr[1] / 2.0
    diff_ex[-1] = -arr[-2] / 2.0
    assert odl_all_equal(diff, diff_ex)

    # boundary: one-sided first-order forward/backward difference without zero
    # padding
    diff = finite_diff(arr, axis=0, method='central', pad_mode='order1')
    diff_ex[0] = arr[1] - arr[0]  # 1st-order accurate forward difference
    diff_ex[-1] = arr[-1] - arr[-2]  # 1st-order accurate backward diff.
    assert odl_all_equal(diff, diff_ex)

    # different edge order really differ
    df1 = finite_diff(arr, axis=0, method='central', pad_mode='order1')
    df2 = finite_diff(arr, axis=0, method='central', pad_mode='order2')
    assert odl_all_equal(df1[1:-1], diff_ex[1:-1])
    assert odl_all_equal(df2[1:-1], diff_ex[1:-1])
    assert df1[0] != df2[0]
    assert df1[-1] != df2[-1]

    # in-place evaluation
    out = ns.zeros_like(arr)
    assert out is finite_diff(arr, axis=0, out=out)
    assert odl_all_equal(out, finite_diff(arr, axis=0))
    assert out is not finite_diff(arr, axis=0)

    # axis
    arr = backend.array_constructor([[0., 2., 4., 6., 8.],
                    [1., 3., 5., 7., 9.]])
    df0 = finite_diff(arr, axis=0, pad_mode='order1')
    darr0 = 1 * ns.ones(arr.shape)
    assert odl_all_equal(df0, darr0)
    darr1 = 2 * ns.ones(arr.shape)
    df1 = finite_diff(arr, axis=1, pad_mode='order1')
    assert odl_all_equal(df1, darr1)

    # complex arrays
    arr = backend.array_constructor([0., 1., 2., 3., 4.]) + 1j * backend.array_constructor([10., 9., 8., 7.,
                                                          6.])
    diff = finite_diff(arr, axis=0, pad_mode='order1')
    assert all(diff.real == 1)
    assert all(diff.imag == -1)


def test_finite_diff_symmetric_padding(data):
    """Finite difference using replicate padding."""

    # Using replicate padding forward and backward differences have zero
    # derivative at the upper or lower endpoint, respectively
    assert finite_diff(data, axis=0, method='forward',
                       pad_mode='symmetric')[-1] == 0
    assert finite_diff(data, axis=0, method='backward',
                       pad_mode='symmetric')[0] == 0

    diff = finite_diff(data, axis=0, method='central', pad_mode='symmetric')
    assert diff[0] == (data[1] - data[0]) / 2
    assert diff[-1] == (data[-1] - data[-2]) / 2


def test_finite_diff_constant_padding(data):
    """Finite difference using constant padding."""

    for pad_const in [-1, 0, 1]:
        diff_forward = finite_diff(data, axis=0, method='forward',
                                   pad_mode='constant',
                                   pad_const=pad_const)

        assert diff_forward[0] == data[1] - data[0]
        assert diff_forward[-1] == pad_const - data[-1]

        diff_backward = finite_diff(data, axis=0, method='backward',
                                    pad_mode='constant',
                                    pad_const=pad_const)

        assert diff_backward[0] == data[0] - pad_const
        assert diff_backward[-1] == data[-1] - data[-2]

        diff_central = finite_diff(data, axis=0, method='central',
                                   pad_mode='constant',
                                   pad_const=pad_const)

        assert diff_central[0] == (data[1] - pad_const) / 2
        assert diff_central[-1] == (pad_const - data[-2]) / 2


def test_finite_diff_periodic_padding(data):
    """Finite difference using periodic padding."""

    diff_forward = finite_diff(data, axis=0, method='forward',
                               pad_mode='periodic')

    assert diff_forward[0] == data[1] - data[0]
    assert diff_forward[-1] == data[0] - data[-1]

    diff_backward = finite_diff(data, axis=0, method='backward',
                                pad_mode='periodic')

    assert diff_backward[0] == data[0] - data[-1]
    assert diff_backward[-1] == data[-1] - data[-2]

    diff_central = finite_diff(data, axis=0, method='central',
                               pad_mode='periodic')

    assert diff_central[0] == (data[1] - data[-1]) / 2
    assert diff_central[-1] == (data[0] - data[-2]) / 2


# --- PartialDerivative --- #


def test_part_deriv_init():
    """Check initialization of ``PartialDerivative``."""
    space = odl.uniform_discr([0, 0], [1, 1], (4, 5))

    op = PartialDerivative(space, axis=0)
    assert repr(op) != ''
    op = PartialDerivative(space, axis=1)
    assert repr(op) != ''
    op = PartialDerivative(space, axis=0, range=space.astype('float32'))
    assert repr(op) != ''
    op = PartialDerivative(space, axis=0, method='central')
    assert repr(op) != ''
    op = PartialDerivative(space, axis=0, pad_const=1)
    assert repr(op) != ''
    op = PartialDerivative(space, axis=0, pad_mode='order1')
    assert repr(op) != ''

    with pytest.raises(TypeError):
        PartialDerivative(odl.rn(1))


def test_part_deriv(space, method, padding):
    """Discretized partial derivative."""
    if isinstance(padding, tuple):
        pad_mode, pad_const = padding
    else:
        pad_mode, pad_const = padding, 0

    dom_vec = noise_element(space)
    dom_vec_arr = dom_vec.asarray()

    for axis in range(space.ndim):
        partial = PartialDerivative(space, axis=axis, method=method,
                                    pad_mode=pad_mode,
                                    pad_const=pad_const)

        # Compare to helper function
        dx = space.cell_sides[axis]
        diff = finite_diff(dom_vec_arr, axis=axis, dx=dx, method=method,
                        pad_mode=pad_mode,
                        pad_const=pad_const)

        partial_vec = partial(dom_vec)
        assert all_almost_equal(partial_vec, diff)

        # Test adjoint operator
        derivative = partial.derivative()
        ran_vec = noise_element(space)
        deriv_vec = derivative(dom_vec)
        adj_vec = derivative.adjoint(ran_vec)
        lhs = ran_vec.inner(deriv_vec)
        rhs = dom_vec.inner(adj_vec)

        # Check not to use trivial data
        assert lhs != 0
        assert rhs != 0
        assert lhs == pytest.approx(rhs, rel=dtype_tol(space.dtype))


# --- Gradient --- #


def test_gradient_init(odl_impl_device_pairs):
    """Check initialization of ``Gradient``."""
    impl, device = odl_impl_device_pairs
    space = odl.uniform_discr([0, 0], [1, 1], (4, 5), impl=impl, device=device)
    vspace = space ** 2

    op = Gradient(space)
    assert repr(op) != ''
    op = Gradient(range=vspace)
    assert repr(op) != ''
    op = Gradient(space, range=space.astype('float32') ** 2)
    assert repr(op) != ''
    op = Gradient(space, method='central')
    assert repr(op) != ''
    op = Gradient(space, pad_const=1)
    assert repr(op) != ''
    op = Gradient(space, pad_mode='order1')
    assert repr(op) != ''

    with pytest.raises(TypeError):
        Gradient(odl.rn(1))

    with pytest.raises(TypeError):
        Gradient(space, range=space)

    with pytest.raises(ValueError):
        Gradient(space, range=space ** 3)


def test_gradient(space, method, padding):
    """Discretized spatial gradient operator."""
    with pytest.raises(TypeError):
        Gradient(odl.rn(1), method=method)

    if isinstance(padding, tuple):
        pad_mode, pad_const = padding
    else:
        pad_mode, pad_const = padding, 0

    # DiscretizedSpaceElement
    dom_vec = noise_element(space)
    dom_vec_arr = dom_vec.asarray()

    # gradient
    grad = Gradient(space, method=method,
                    pad_mode=pad_mode,
                    pad_const=pad_const)
    grad_vec = grad(dom_vec)
    assert len(grad_vec) == space.ndim

    # computation of gradient components with helper function
    for axis, dx in enumerate(space.cell_sides):
        diff = finite_diff(dom_vec_arr, axis=axis, dx=dx, method=method,
                        pad_mode=pad_mode,
                        pad_const=pad_const)

        assert all_almost_equal(grad_vec[axis].asarray(), diff)

    # Test adjoint operator
    derivative = grad.derivative()
    ran_vec = noise_element(derivative.range)
    deriv_grad_vec = derivative(dom_vec)
    adj_grad_vec = derivative.adjoint(ran_vec)
    lhs = ran_vec.inner(deriv_grad_vec)
    rhs = dom_vec.inner(adj_grad_vec)

    # Check not to use trivial data
    assert lhs != 0
    assert rhs != 0
    assert lhs == pytest.approx(rhs, rel=dtype_tol(space.dtype))

    # Higher-dimensional arrays
    lin_size = 3
    for ndim in [1, 3, 6]:
        space = odl.uniform_discr([0.] * ndim, [1.] * ndim, [lin_size] * ndim)
        dom_vec = odl.core.phantom.cuboid(space, [0.2] * ndim, [0.8] * ndim)

        grad = Gradient(space, method=method, pad_mode=pad_mode,
                        pad_const=pad_const)
        grad(dom_vec)

# --- Divergence --- #


def test_divergence_init(odl_impl_device_pairs):
    """Check initialization of ``Divergence``."""
    impl, device = odl_impl_device_pairs
    space = odl.uniform_discr([0, 0], [1, 1], (4, 5), impl=impl, device=device)
    vspace = space ** 2

    op = Divergence(vspace)
    assert repr(op) != ''
    op = Divergence(range=space)
    assert repr(op) != ''
    op = Divergence(vspace, range=space.astype('float32'))
    assert repr(op) != ''
    op = Divergence(vspace, method='central')
    assert repr(op) != ''
    op = Divergence(vspace, pad_const=1)
    assert repr(op) != ''
    op = Divergence(vspace, pad_mode='order1')
    assert repr(op) != ''

    with pytest.raises(TypeError):
        Divergence(odl.rn(1) ** 2)

    with pytest.raises(TypeError):
        Divergence(vspace, range=vspace)

    with pytest.raises(ValueError):
        Divergence(space ** 3, range=space)


def test_divergence(space, method, padding):
    """Discretized spatial divergence operator."""
    # Invalid space
    with pytest.raises(TypeError):
        Divergence(range=odl.rn(1), method=method)

    if isinstance(padding, tuple):
        pad_mode, pad_const = padding
    else:
        pad_mode, pad_const = padding, 0

    # Operator instance
    div = Divergence(range=space, method=method,
                    pad_mode=pad_mode,
                    pad_const=pad_const)

    # Apply operator
    dom_vec = noise_element(div.domain)
    div_dom_vec = div(dom_vec)

    # computation of divergence with helper function
    expected_result = space.array_namespace.zeros(space.shape, dtype=space.dtype, device=space.device)
    for axis, dx in enumerate(space.cell_sides):
        expected_result += finite_diff(dom_vec[axis], axis=axis, dx=dx,
                                    method=method, pad_mode=pad_mode,
                                    pad_const=pad_const)

    assert all_almost_equal(expected_result, div_dom_vec.asarray())

    # Adjoint operator
    derivative = div.derivative()
    deriv_div_dom_vec = derivative(dom_vec)
    ran_vec = noise_element(div.range)
    adj_div_ran_vec = derivative.adjoint(ran_vec)

    # Adjoint condition
    lhs = ran_vec.inner(deriv_div_dom_vec)
    rhs = dom_vec.inner(adj_div_ran_vec)
    # Check not to use trivial data
    assert lhs != 0
    assert rhs != 0
    assert lhs == pytest.approx(rhs, rel=dtype_tol(space.dtype))


# --- Laplacian --- #

def test_laplacian_init(odl_impl_device_pairs):
    """Check initialization of ``Laplacian``."""
    impl, device = odl_impl_device_pairs
    space = odl.uniform_discr([0, 0], [1, 1], (4, 5), impl=impl, device=device)

    op = Laplacian(space)
    assert repr(op) != ''
    op = Laplacian(space, range=space.astype('float32'))
    assert repr(op) != ''
    op = Laplacian(space, pad_const=1)
    assert repr(op) != ''
    op = Laplacian(space, pad_mode='order0')
    assert repr(op) != ''

    with pytest.raises(TypeError):
        Laplacian(odl.rn(1))


def test_laplacian(space, padding):
    """Discretized spatial laplacian operator."""
    if isinstance(padding, tuple):
        pad_mode, pad_const = padding
    else:
        pad_mode, pad_const = padding, 0

    if pad_mode in ('order1', 'order2'):
        return  # these pad modes not supported for laplacian

    # Operator instance
    lap = Laplacian(space, pad_mode=pad_mode, pad_const=pad_const)

    # Apply operator
    dom_vec = noise_element(space)
    div_dom_vec = lap(dom_vec)

    # computation of divergence with helper function
    expected_result = space.array_namespace.zeros(space.shape, device=space.device, dtype=space.dtype)
    for axis, dx in enumerate(space.cell_sides):
        diff_f = finite_diff(dom_vec.asarray(), axis=axis, dx=dx ** 2,
                            method='forward', pad_mode=pad_mode,
                            pad_const=pad_const)
        diff_b = finite_diff(dom_vec.asarray(), axis=axis, dx=dx ** 2,
                            method='backward', pad_mode=pad_mode,
                            pad_const=pad_const)
        expected_result += diff_f - diff_b

    assert all_almost_equal(expected_result, div_dom_vec.asarray())

    # Adjoint operator
    derivative = lap.derivative()
    deriv_lap_dom_vec = derivative(dom_vec)
    ran_vec = noise_element(lap.range)
    adj_lap_ran_vec = derivative.adjoint(ran_vec)

    # Adjoint condition
    lhs = ran_vec.inner(deriv_lap_dom_vec)
    rhs = dom_vec.inner(adj_lap_ran_vec)

    # Check not to use trivial data
    assert lhs != 0
    assert rhs != 0
    assert lhs == pytest.approx(rhs, rel=dtype_tol(space.dtype))


if __name__ == '__main__':
    odl.core.util.test_file(__file__)
