# 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/.

"""Tests for linearized deformation operators."""

from __future__ import division

import numpy as np
import pytest

import odl
from odl.trafos.deform import LinDeformFixedDisp, LinDeformFixedTempl
from odl.core.util.testutils import simple_fixture

from odl.core.array_API_support import get_array_and_backend, exp

# --- pytest fixtures --- #


dtype = simple_fixture('dtype', [float, complex])
interp = simple_fixture('interp', ['linear', 'nearest'])
ndim = simple_fixture('ndim', [1, 2, 3])


@pytest.fixture
def space(ndim, dtype, odl_impl_device_pairs):
    """Provide a space for unit tests."""
    impl, device = odl_impl_device_pairs
    # supported_dtypes = odl.lookup_array_backend(impl).available_dtypes
    # # if np.dtype(dtype) not in supported_dtypes:
    # #     pytest.skip('dtype not available for this backend')

    return odl.uniform_discr(
        [-1] * ndim, [1] * ndim, [20] * ndim, impl=impl, dtype=dtype, device=device
    )


# --- Helper functions --- #


SIGMA = 0.3  # width of the gaussian
EPS = 0.25   # scale of the displacement field


def error_bound(interp):
    """Error bound varies with interpolation (larger for "worse")."""
    if interp == 'linear':
        return 0.1
    elif interp == 'nearest':
        return 0.2


def prod(x):
    """Product of a sequence."""
    prod = 1
    for xi in x:
        prod = prod * xi
    return prod


def numpy_template_function(x):
    """Gaussian function with std SIGMA."""
    return np.exp(-sum(xi ** 2 for xi in x) / SIGMA ** 2)

def torch_template_function(x):
    """Gaussian function with std SIGMA."""
    import torch
    return torch.exp(-sum(xi ** 2 for xi in x) / SIGMA ** 2)

def numpy_template_grad_factory(n):
    """Gradient of the gaussian."""
    def template_grad_i(i):
        # Indirection for lambda capture
        return lambda x: -2 * x[i] / SIGMA ** 2 * numpy_template_function(x)
    return [template_grad_i(i) for i in range(n)]

def torch_template_grad_factory(n):
    """Gradient of the gaussian."""
    def template_grad_i(i):
        # Indirection for lambda capture
        return lambda x: -2 * x[i] / SIGMA ** 2 * torch_template_function(x)
    return [template_grad_i(i) for i in range(n)]


def disp_field_factory(n):
    """Displacement field.

    In 1d: (x)
    In 2d: (xy, y)
    In 3d: (xyz, y, z)
    etc...
    """
    def coordinate_projection_i(i):
        # Indirection for lambda capture
        return lambda x: EPS * x[i]

    lst = [lambda x: EPS * prod(x)]
    lst.extend(coordinate_projection_i(i) for i in range(1, n))
    return lst


def numpy_exp_div_inv_disp(x):
    """Exponential of the divergence of the displacement field.

    In 1d: exp(- EPS)
    In 2d: exp(- EPS * (y + 1))
    In 2d: exp(- EPS * (yz + 2))
    """
    return np.exp(- EPS * (prod(x[1:]) + (len(x) - 1)))


def torch_exp_div_inv_disp(x):
    """Exponential of the divergence of the displacement field.

    In 1d: exp(- EPS)
    In 2d: exp(- EPS * (y + 1))
    In 2d: exp(- EPS * (yz + 2))
    """
    import torch
    return torch.exp(- EPS * (prod(x[1:]) + (len(x) - 1)))


def displaced_points(x):
    """Displaced coordinate points."""
    disp = [dsp(x) for dsp in disp_field_factory(len(x))]
    return [xi + di for xi, di in zip(x, disp)]


def numpy_deformed_template(x):
    """Deformed template."""
    return numpy_template_function(displaced_points(x))

def torch_deformed_template(x):
    """Deformed template."""
    return torch_template_function(displaced_points(x))

def vector_field_factory(n):
    """Vector field for the gradient.

    In 1d: (x)
    In 2d: (x, y)
    In 3d: (x, y, z)
    etc...
    """
    def vector_field_i(i):
        return lambda x: x[i]
    return [vector_field_i(i) for i in range(n)]


def numpy_template_deformed_grad_factory(n):
    """Deformed gradient."""
    templ_grad = numpy_template_grad_factory(n)

    def template_deformed_gradi(i):
        # Indirection for lambda capture
        return lambda x: templ_grad[i](displaced_points(x))

    return [template_deformed_gradi(i) for i in range(n)]

def torch_template_deformed_grad_factory(n):
    """Deformed gradient."""
    templ_grad = torch_template_grad_factory(n)

    def template_deformed_gradi(i):
        # Indirection for lambda capture
        return lambda x: templ_grad[i](displaced_points(x))

    return [template_deformed_gradi(i) for i in range(n)]

def numpy_fixed_templ_deriv(x):
    """Derivative taken in disp_field and evaluated in vector_field."""
    dg = [tdgf(x) for tdgf in numpy_template_deformed_grad_factory(len(x))]
    v = [vff(x) for vff in vector_field_factory(len(x))]
    return sum(dgi * vi for dgi, vi in zip(dg, v))

def torch_fixed_templ_deriv(x):
    """Derivative taken in disp_field and evaluated in vector_field."""
    dg = [tdgf(x) for tdgf in torch_template_deformed_grad_factory(len(x))]
    v = [vff(x) for vff in vector_field_factory(len(x))]
    return sum(dgi * vi for dgi, vi in zip(dg, v))


def numpy_inv_deformed_template(x):
    """Analytic inverse deformation of the template function."""
    disp = [dsp(x) for dsp in disp_field_factory(len(x))]
    disp_x = [xi - di for xi, di in zip(x, disp)]
    return numpy_template_function(disp_x)

def torch_inv_deformed_template(x):
    """Analytic inverse deformation of the template function."""
    disp = [dsp(x) for dsp in disp_field_factory(len(x))]
    disp_x = [xi - di for xi, di in zip(x, disp)]
    return torch_template_function(disp_x)


# --- LinDeformFixedTempl --- #


def test_fixed_templ_init(odl_impl_device_pairs):
    """Test init and props of linearized deformation with fixed template."""
    impl, device = odl_impl_device_pairs
    space = odl.uniform_discr(0, 1, 5, impl=impl,device=device)

    if impl == 'numpy':
        template_function = numpy_template_function 
    else:
        template_function = torch_template_function

    template = space.element(template_function)

    # Valid input
    op = LinDeformFixedTempl(template)
    assert repr(op) != ''
    op = LinDeformFixedTempl(template, domain=space.astype('float32') ** 1)
    assert repr(op) != ''

    # Invalid input
    with pytest.raises(TypeError):
        # template_function not a DiscretizedSpaceElement
        LinDeformFixedTempl(template_function)

@pytest.fixture
def space(odl_impl_device_pairs):
    """Provide a space for unit tests."""
    impl, device = odl_impl_device_pairs
    ndim = 2
    # supported_dtypes = odl.lookup_array_backend(impl).available_dtypes
    # # if np.dtype(dtype) not in supported_dtypes:
    # #     pytest.skip('dtype not available for this backend')

    return odl.uniform_discr(
        [-1] * ndim, [1] * ndim, [20] * ndim, impl=impl, device=device
    )


def test_fixed_templ_call(space, interp):
    """Test call of linearized deformation with fixed template."""
    # Define the analytic template as the hat function and its gradient
    if space.impl == 'numpy':
        template_function = numpy_template_function 
        deformed_template = numpy_deformed_template
    else:
        if space.dtype_identifier == 'complex128':
            return 
        template_function = torch_template_function
        deformed_template = torch_deformed_template

    template = space.element(template_function)
    deform_op = LinDeformFixedTempl(template, interp=interp)

    # Calculate result and exact result
    true_deformed_templ = space.element(deformed_template)

    deformed_templ = deform_op(disp_field_factory(space.ndim))

    # Verify that the result is within error limits
    error = (true_deformed_templ - deformed_templ).norm()
    rlt_err = error / deformed_templ.norm()
    assert rlt_err < error_bound(interp)


def test_fixed_templ_deriv(space, interp):
    """Test derivative of linearized deformation with fixed template."""
    if not space.is_real:
        pytest.skip('derivative not implemented for complex dtypes')

    if space.impl == 'numpy':
        template_function = numpy_template_function 
        fixed_templ_deriv = numpy_fixed_templ_deriv
    else:
        template_function = torch_template_function
        fixed_templ_deriv = torch_fixed_templ_deriv

    # Set up template and displacement field
    template = space.element(template_function)
    disp_field = disp_field_factory(space.ndim)
    vector_field = vector_field_factory(space.ndim)
    fixed_templ_op = LinDeformFixedTempl(template, interp=interp)

    # Calculate result
    fixed_templ_op_deriv = fixed_templ_op.derivative(disp_field)
    fixed_templ_deriv_comp = fixed_templ_op_deriv(vector_field)

    # Calculate the analytic result
    fixed_templ_deriv_exact = space.element(fixed_templ_deriv)

    # Verify that the result is within error limits
    error = (fixed_templ_deriv_exact - fixed_templ_deriv_comp).norm()
    rlt_err = error / fixed_templ_deriv_comp.norm()
    assert rlt_err < error_bound(interp)


# --- LinDeformFixedDisp --- #


def test_fixed_disp_init(odl_impl_device_pairs):
    """Test init and props of lin. deformation with fixed displacement."""
    impl, device=odl_impl_device_pairs
    space = odl.uniform_discr(0, 1, 5, impl=impl, device=device)
    disp_field = space.tangent_bundle.element(
        disp_field_factory(space.ndim))

    # Valid input
    op = LinDeformFixedDisp(disp_field)
    assert repr(op) != ''
    op = LinDeformFixedDisp(disp_field, templ_space=space)
    assert repr(op) != ''

    # Non-valid input
    with pytest.raises(TypeError):  # displacement not ProductSpaceElement
        LinDeformFixedDisp(space.one())
    with pytest.raises(TypeError):  # templ_space not DiscretizedSpace
        LinDeformFixedDisp(disp_field, space.tangent_bundle)
    with pytest.raises(TypeError):  # templ_space not a power space
        bad_pspace = odl.ProductSpace(space, odl.rn(3))
        LinDeformFixedDisp(disp_field, bad_pspace)
    with pytest.raises(TypeError):  # templ_space not based on DiscretizedSpace
        bad_pspace = odl.ProductSpace(odl.rn(2), 1)
        LinDeformFixedDisp(disp_field, bad_pspace)
    with pytest.raises(TypeError):  # wrong dtype on templ_space
        wrong_dtype = odl.ProductSpace(space.astype(complex), 1)
        LinDeformFixedDisp(disp_field, wrong_dtype)
    with pytest.raises(ValueError):  # vector field spaces don't match
        bad_space = odl.uniform_discr(0, 1, 10)
        LinDeformFixedDisp(disp_field, bad_space)


def test_fixed_disp_call(space, interp):
    """Test call of lin. deformation with fixed displacement."""
    if space.impl == 'numpy':
        template_function = numpy_template_function 
        deformed_template = numpy_deformed_template
    else:
        if space.dtype_identifier == 'complex128':
            return 
        template_function = torch_template_function
        deformed_template = torch_deformed_template

    template = space.element(template_function)
    disp_field = space.real_space.tangent_bundle.element(
        disp_field_factory(space.ndim))

    # Calculate result and exact result
    deform_op = LinDeformFixedDisp(
        disp_field, templ_space=space, interp=interp
    )
    deformed_templ = deform_op(template)
    true_deformed_templ = space.element(deformed_template)

    # Verify that the result is within error limits
    error = (true_deformed_templ - deformed_templ).norm()
    rlt_err = error / deformed_templ.norm()
    assert rlt_err < error_bound(interp)


def test_fixed_disp_inv(space, interp):
    """Test inverse of lin. deformation with fixed displacement."""
    if space.impl == 'numpy':
        template_function = numpy_template_function 
    else:
        if space.dtype_identifier == 'complex128':
            return 
        template_function = torch_template_function
        
    # Set up template and displacement field
    template = space.element(template_function)
    disp_field = space.real_space.tangent_bundle.element(
        disp_field_factory(space.ndim))

    # Verify that the inverse is in fact a (left and right) inverse
    deform_op = LinDeformFixedDisp(
        disp_field, templ_space=space, interp=interp
    )

    result_op_inv = deform_op(deform_op.inverse(template))
    error = (result_op_inv - template).norm()
    rel_err = error / template.norm()
    assert rel_err < 2 * error_bound(interp)  # need a bit more tolerance

    result_inv_op = deform_op.inverse(deform_op(template))
    error = (result_inv_op - template).norm()
    rel_err = error / template.norm()
    assert rel_err < 2 * error_bound(interp)  # need a bit more tolerance


def test_fixed_disp_adj(space, interp):
    """Test adjoint of lin. deformation with fixed displacement."""
    # Set up template and displacement field

    if space.impl == 'numpy':
        template_function = numpy_template_function 
        inv_deformed_template = numpy_inv_deformed_template
        exp_div_inv_disp = numpy_exp_div_inv_disp
    else:
        if space.dtype_identifier == 'complex128':
            return 
        template_function = torch_template_function
        inv_deformed_template = torch_inv_deformed_template
        exp_div_inv_disp = torch_exp_div_inv_disp

    template = space.element(template_function)
    disp_field = space.real_space.tangent_bundle.element(
        disp_field_factory(space.ndim))

    # Calculate result
    deform_op = LinDeformFixedDisp(
        disp_field, templ_space=space, interp=interp
    )
    deformed_templ_adj = deform_op.adjoint(template)

    # Calculate the analytic result
    true_deformed_templ_adj = space.element(inv_deformed_template)
    exp_div = space.element(exp_div_inv_disp)
    true_deformed_templ_adj *= exp_div

    # Verify that the result is within error limits
    error = (deformed_templ_adj - true_deformed_templ_adj).norm()
    rel_err = error / true_deformed_templ_adj.norm()
    assert rel_err < error_bound(interp)

    # Verify the adjoint definition <Ax, x> = <x, A^* x>
    deformed_templ = deform_op(template)
    inner1 = deformed_templ.inner(template)
    inner2 = template.inner(deformed_templ_adj)
    assert inner1 == pytest.approx(inner2, abs=.1)


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