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

from __future__ import division

import numpy as np
import pytest

import odl
from odl.trafos.fourier import (
    DiscreteFourierTransform, DiscreteFourierTransformInverse,
    FourierTransform)
from odl.trafos.util.ft_utils import (
    _interp_kernel_ft, dft_postprocess_data, dft_preprocess_data,
    reciprocal_grid)
from odl.core.util import (
    all_almost_equal, complex_dtype, conj_exponent, is_real_dtype,
    noise_element, skip_if_no_pyfftw)
from odl.core.util.testutils import simple_fixture

from odl.core.array_API_support import allclose

# --- pytest fixtures --- #


impl = simple_fixture(
    'impl',
    [pytest.param('numpy'),
     pytest.param('pyfftw', marks=skip_if_no_pyfftw),
     pytest.param('default')]
)
exponent = simple_fixture('exponent', [2.0, 1.0, float('inf'), 1.5])
sign = simple_fixture('sign', ['-', '+'])


# --- Helper functions --- #


def _params_from_dtype(dtype):
    if is_real_dtype(dtype):
        halfcomplex = True
    else:
        halfcomplex = False
    return halfcomplex, complex_dtype(dtype)


@pytest.fixture
def _dft_complex_space(odl_complex_floating_dtype, odl_impl_device_pairs):
    impl, device = odl_impl_device_pairs
    shape = (4,5)
    return odl.uniform_discr(
        [0] * 2,
        np.subtract(shape, 1),
        shape,
        dtype=odl_complex_floating_dtype,
        nodes_on_bdry=True,
        impl=impl, device=device
    )
def skip_incompatible_impl(impl, dft_space):
    if dft_space.impl != 'numpy' and impl!= 'default':
        pytest.skip('Currently, only the numpy backend supports fft implementations that are not array-API compatible (e.g pyfftw)')

def _dft_space(shape, dtype='float64'):
    try:
        ndim = len(shape)
    except TypeError:
        ndim = 1
    return odl.uniform_discr(
        [0] * ndim,
        np.subtract(shape, 1),
        shape,
        dtype=dtype,
        nodes_on_bdry=True,
    )


def sinc(x):
    # numpy.sinc scales by pi, we don't want that
    return np.sinc(x / np.pi)

def sinc_pytorch(x):
    # torch.sinc scales by pi, we don't want that
    import torch 
    return torch.sinc(x / np.pi)


# ---- DiscreteFourierTransform ---- #


def test_dft_init(impl, odl_impl_device_pairs):
    # Just check if the code runs at all
    backend_impl, device = odl_impl_device_pairs
    shape = (4, 5)
    dom = _dft_space(shape)
    dom_nonseq = odl.uniform_discr([0, 0], [1, 1], shape, impl=backend_impl, device=device)
    dom_f32 = dom.astype('float32')
    ran = _dft_space(shape, dtype='complex128')
    ran_c64 = ran.astype('complex64')
    ran_hc = _dft_space((3, 5), dtype='complex128')

    # Implicit range
    DiscreteFourierTransform(dom, impl=impl)
    DiscreteFourierTransform(dom_nonseq, impl=impl)
    DiscreteFourierTransform(dom_f32, impl=impl)
    DiscreteFourierTransform(dom, axes=(0,), impl=impl)
    DiscreteFourierTransform(dom, axes=(0, -1), impl=impl)
    DiscreteFourierTransform(dom, axes=(0,), halfcomplex=True, impl=impl)
    DiscreteFourierTransform(dom, impl=impl, sign='+')

    # Explicit range
    DiscreteFourierTransform(dom, range=ran, impl=impl)
    DiscreteFourierTransform(dom_f32, range=ran_c64, impl=impl)
    DiscreteFourierTransform(dom, range=ran, axes=(0,), impl=impl)
    DiscreteFourierTransform(dom, range=ran, axes=(0,), impl=impl, sign='+')
    DiscreteFourierTransform(dom, range=ran, axes=(0, -1), impl=impl)
    DiscreteFourierTransform(dom, range=ran_hc, axes=(0,), impl=impl,
                             halfcomplex=True)


def test_dft_init_raise():
    # Test different error scenarios
    shape = (4, 5)
    dom = _dft_space(shape)
    dom_f32 = _dft_space(shape, dtype='float32')

    # Bad types
    with pytest.raises(TypeError):
        DiscreteFourierTransform(dom.tspace)

    with pytest.raises(TypeError):
        DiscreteFourierTransform(dom, dom.tspace)

    # Illegal arguments
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, impl='fftw')

    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, axes=(1, 2))

    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, axes=(1, -3))

    # Badly shaped range
    bad_ran = _dft_space((3, 5), dtype='complex128')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran)

    bad_ran = _dft_space((10, 10), dtype='complex128')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran)

    bad_ran = _dft_space((4, 5), dtype='complex128')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran, halfcomplex=True)

    bad_ran = _dft_space((4, 3), dtype='complex128')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran, halfcomplex=True, axes=(0,))

    # Bad data types
    bad_ran = _dft_space(shape, dtype='complex64')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran)

    bad_ran = _dft_space(shape, dtype='float64')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran)

    bad_ran = _dft_space((4, 3), dtype='float64')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, bad_ran, halfcomplex=True)

    bad_ran = _dft_space((4, 3), dtype='complex128')
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom_f32, bad_ran, halfcomplex=True)

    # Bad sign
    with pytest.raises(ValueError):
        DiscreteFourierTransform(dom, sign=-1)


def test_dft_range():
    # 1d
    shape = 10
    dom = _dft_space(shape, dtype='complex128')
    fft = DiscreteFourierTransform(dom)
    true_ran = _dft_space(shape, dtype='complex128')
    assert fft.range == true_ran

    # 3d
    shape = (3, 4, 5)
    ran = _dft_space(shape, dtype='complex64')
    fft = DiscreteFourierTransform(ran)
    true_ran = _dft_space(shape, dtype='complex64')
    assert fft.range == true_ran

    # 3d, with axes and halfcomplex
    shape = (3, 4, 5)
    axes = (-1, -2)
    ran_shape = (3, 3, 5)
    dom = _dft_space(shape, dtype='float32')
    fft = DiscreteFourierTransform(dom, axes=axes, halfcomplex=True)
    true_ran = _dft_space(ran_shape, dtype='complex64')
    assert fft.range == true_ran


# ---- DiscreteFourierTransformInverse ---- #


def test_idft_init(impl):
    # Just check if the code runs at all; this uses the init function of
    # DiscreteFourierTransform, so we don't need exhaustive tests here
    shape = (4, 5)
    ran = _dft_space(shape, dtype='complex128')
    ran_hc = _dft_space(shape, dtype='float64')
    dom = _dft_space(shape, dtype='complex128')
    dom_hc = _dft_space((3, 5), dtype='complex128')

    # Implicit range
    DiscreteFourierTransformInverse(dom, impl=impl)

    # Explicit range
    DiscreteFourierTransformInverse(ran, domain=dom, impl=impl)
    DiscreteFourierTransformInverse(ran_hc, domain=dom_hc, axes=(0,),
                                    impl=impl, halfcomplex=True)


def test_dft_call(impl, _dft_complex_space):
    # 2d, complex, all ones and random back & forth
    # dft_dom = _dft_space(shape, dtype='complex64')
    skip_incompatible_impl(impl,_dft_complex_space)
    dft_dom = _dft_complex_space
    dft = DiscreteFourierTransform(domain=dft_dom, impl=impl)
    idft = DiscreteFourierTransformInverse(range=dft_dom, impl=impl)

    assert dft.domain == idft.range
    assert dft.range == idft.domain

    one = dft.domain.one()
    one_dft1 = dft(one, flags=('FFTW_ESTIMATE',))
    one_dft2 = dft.inverse.inverse(one, flags=('FFTW_ESTIMATE',))
    one_dft3 = dft.adjoint.adjoint(one, flags=('FFTW_ESTIMATE',))
    true_dft = [[20, 0, 0, 0, 0],  # along all axes by default
                [0, 0, 0, 0, 0],
                [0, 0, 0, 0, 0],
                [0, 0, 0, 0, 0]]
    assert all_almost_equal(one_dft1, true_dft)
    assert all_almost_equal(one_dft2, true_dft)
    assert all_almost_equal(one_dft3, true_dft)

    one_idft1 = idft(one_dft1, flags=('FFTW_ESTIMATE',))
    one_idft2 = dft.inverse(one_dft1, flags=('FFTW_ESTIMATE',))
    one_idft3 = dft.adjoint(one_dft1, flags=('FFTW_ESTIMATE',))
    assert all_almost_equal(one_idft1, one)
    assert all_almost_equal(one_idft2, one)
    assert all_almost_equal(one_idft3, one)

    rand_arr = noise_element(dft_dom)
    rand_arr_dft = dft(rand_arr, flags=('FFTW_ESTIMATE',))
    rand_arr_idft = idft(rand_arr_dft, flags=('FFTW_ESTIMATE',))
    assert (rand_arr_idft - rand_arr).norm() < 1e-6

    # 2d, halfcomplex, first axis
    axes = 0
    dft_dom = _dft_complex_space.real_space
    dft = DiscreteFourierTransform(domain=dft_dom, impl=impl, halfcomplex=True,
                                   axes=axes)
    idft = DiscreteFourierTransformInverse(range=dft_dom, impl=impl,
                                           halfcomplex=True, axes=axes)

    assert dft.domain == idft.range
    assert dft.range == idft.domain

    one = dft.domain.one()
    one_dft = dft(one, flags=('FFTW_ESTIMATE',))
    true_dft = [[4, 4, 4, 4, 4],  # transform axis shortened
                [0, 0, 0, 0, 0],
                [0, 0, 0, 0, 0]]
    assert all_almost_equal(one_dft, true_dft)

    one_idft1 = idft(one_dft, flags=('FFTW_ESTIMATE',))
    one_idft2 = dft.inverse(one_dft, flags=('FFTW_ESTIMATE',))
    assert all_almost_equal(one_idft1, one)
    assert all_almost_equal(one_idft2, one)

    rand_arr = noise_element(dft_dom)
    rand_arr_dft = dft(rand_arr, flags=('FFTW_ESTIMATE',))
    rand_arr_idft = idft(rand_arr_dft, flags=('FFTW_ESTIMATE',))
    assert (rand_arr_idft - rand_arr).norm() < 1e-6


def test_dft_sign(impl, _dft_complex_space):
    # Test if the FT sign behaves as expected, i.e. that the FT with sign
    # '+' and '-' have same real parts and opposite imaginary parts.
    # 2d, complex, all ones and random back & forth
    skip_incompatible_impl(impl, _dft_complex_space)
    dft_dom = _dft_complex_space
    dft_minus = DiscreteFourierTransform(domain=dft_dom, impl=impl, sign='-')
    dft_plus = DiscreteFourierTransform(domain=dft_dom, impl=impl, sign='+')

    arr = dft_dom.element([[0, 0, 0, 0, 0],
                           [0, 0, 1, 1, 0],
                           [0, 0, 1, 1, 0],
                           [0, 0, 0, 0, 0]])
    arr_dft_minus = dft_minus(arr, flags=('FFTW_ESTIMATE',))
    arr_dft_plus = dft_plus(arr, flags=('FFTW_ESTIMATE',))

    assert all_almost_equal(arr_dft_minus.real, arr_dft_plus.real)
    assert all_almost_equal(arr_dft_minus.imag, -arr_dft_plus.imag)
    assert all_almost_equal(dft_minus.inverse(arr_dft_minus), arr)
    assert all_almost_equal(dft_plus.inverse(arr_dft_plus), arr)
    assert all_almost_equal(dft_minus.inverse.inverse(arr), dft_minus(arr))
    assert all_almost_equal(dft_plus.inverse.inverse(arr), dft_plus(arr))

    # 2d, halfcomplex, first axis
    axes = (0,)
    dft_dom = _dft_complex_space.real_space
    arr = dft_dom.element([[0, 0, 0, 0, 0],
                           [0, 0, 1, 1, 0],
                           [0, 0, 1, 1, 0],
                           [0, 0, 0, 0, 0]])

    dft = DiscreteFourierTransform(
        domain=dft_dom, impl=impl, halfcomplex=True, sign='-', axes=axes)

    arr_dft_minus = dft(arr, flags=('FFTW_ESTIMATE',))
    arr_idft_minus = dft.inverse(arr_dft_minus, flags=('FFTW_ESTIMATE',))

    assert all_almost_equal(arr_idft_minus, arr)

    with pytest.raises(ValueError):
        DiscreteFourierTransform(
            domain=dft_dom, impl=impl, halfcomplex=True, sign='+', axes=axes)


def test_dft_init_plan(impl):

    # 2d, halfcomplex, first axis
    shape = (4, 5)
    axes = 0
    dft_dom = _dft_space(shape, dtype='float32')

    dft = DiscreteFourierTransform(dft_dom, impl=impl, axes=axes,
                                   halfcomplex=True)

    if impl != 'pyfftw':
        with pytest.raises(ValueError):
            dft.init_fftw_plan()
        with pytest.raises(ValueError):
            dft.clear_fftw_plan()
    else:
        dft.init_fftw_plan()

        # Make sure plan can be used
        dft._fftw_plan(dft.domain.element().asarray(),
                       dft.range.element().asarray())
        dft.clear_fftw_plan()
        assert dft._fftw_plan is None


# ---- FourierTransform ---- #


def test_fourier_trafo_range(exponent, odl_floating_dtype, odl_impl_device_pairs):
    # Check if the range is initialized correctly. Encompasses the init test
    dtype = odl_floating_dtype

    # Testing R2C for real dtype, else C2C
    impl, device = odl_impl_device_pairs

    # 1D
    shape = 10
    space_discr = odl.uniform_discr(0, 1, shape, exponent=exponent,
                                    impl=impl, device=device, dtype=dtype)

    dft = FourierTransform(space_discr, halfcomplex=True, shift=True)
    assert dft.range.field == odl.ComplexNumbers()
    halfcomplex = True if is_real_dtype(dtype) else False
    assert dft.range.grid == reciprocal_grid(dft.domain.grid,
                                             halfcomplex=halfcomplex,
                                             shift=True)
    assert dft.range.exponent == conj_exponent(exponent)

    # 3D
    shape = (3, 4, 5)
    space_discr = odl.uniform_discr([0] * 3, [1] * 3, shape, exponent=exponent,
                                    impl=impl, device=device, dtype=dtype)

    dft = FourierTransform(space_discr, halfcomplex=True, shift=True)
    assert dft.range.field == odl.ComplexNumbers()
    halfcomplex = True if is_real_dtype(dtype) else False
    assert dft.range.grid == reciprocal_grid(dft.domain.grid,
                                             halfcomplex=halfcomplex,
                                             shift=True)
    assert dft.range.exponent == conj_exponent(exponent)

    # shift must be True in the last axis
    if halfcomplex:
        with pytest.raises(ValueError):
            FourierTransform(space_discr, shift=(True, True, False))

    if exponent != 2.0:
        with pytest.raises(NotImplementedError):
            dft.adjoint

    with pytest.raises(TypeError):
        FourierTransform(dft.domain.partition)


def test_fourier_trafo_init_plan(impl, odl_floating_dtype):
    dtype = odl_floating_dtype

    # Not supported, skip
    if dtype == np.dtype('float16') and impl == 'pyfftw':
        return

    shape = 10
    halfcomplex, _ = _params_from_dtype(dtype)

    space_discr = odl.uniform_discr(0, 1, shape, dtype=dtype)

    ft = FourierTransform(space_discr, impl=impl, halfcomplex=halfcomplex)
    if impl != 'pyfftw':
        with pytest.raises(ValueError):
            ft.init_fftw_plan()
        with pytest.raises(ValueError):
            ft.clear_fftw_plan()
    else:
        ft.init_fftw_plan()

        # Make sure plan can be used
        ft._fftw_plan(ft.domain.element().asarray(),
                      ft.range.element().asarray())
        ft.clear_fftw_plan()
        assert ft._fftw_plan is None

    # With temporaries
    ft.create_temporaries(r=True, f=False)
    if impl != 'pyfftw':
        with pytest.raises(ValueError):
            ft.init_fftw_plan()
        with pytest.raises(ValueError):
            ft.clear_fftw_plan()
    else:
        ft.init_fftw_plan()

        # Make sure plan can be used
        ft._fftw_plan(ft.domain.element().asarray(),
                      ft.range.element().asarray())
        ft.clear_fftw_plan()
        assert ft._fftw_plan is None

    ft.create_temporaries(r=False, f=True)
    if impl != 'pyfftw':
        with pytest.raises(ValueError):
            ft.init_fftw_plan()
        with pytest.raises(ValueError):
            ft.clear_fftw_plan()
    else:
        ft.init_fftw_plan()

        # Make sure plan can be used
        ft._fftw_plan(ft.domain.element().asarray(),
                      ft.range.element().asarray())
        ft.clear_fftw_plan()
        assert ft._fftw_plan is None


def test_fourier_trafo_create_temp():

    shape = 10
    space_discr = odl.uniform_discr(0, 1, shape, dtype='complex64')

    ft = FourierTransform(space_discr)
    ft.create_temporaries()
    assert ft._tmp_r is not None
    assert ft._tmp_f is not None

    ift = ft.inverse
    assert ift._tmp_r is not None
    assert ift._tmp_f is not None

    ft.clear_temporaries()
    assert ft._tmp_r is None
    assert ft._tmp_f is None


def test_fourier_trafo_call(impl, odl_floating_dtype, odl_impl_device_pairs):
    # Test if all variants can be called without error
    backend_impl, device = odl_impl_device_pairs
    dtype = odl_floating_dtype

    # Not supported, skip
    if dtype == np.dtype('float16') and impl == 'pyfftw':
        return

    shape = 10
    halfcomplex, _ = _params_from_dtype(dtype)
    space_discr = odl.uniform_discr(0, 1, shape, dtype=dtype, impl=backend_impl, device=device)

    skip_incompatible_impl(impl, space_discr)

    ft = FourierTransform(space_discr, impl=impl, halfcomplex=halfcomplex)
    ift = ft.inverse

    one = space_discr.one()
    assert odl.allclose(ift(ft(one)), one)

    # With temporaries
    ft.create_temporaries()
    ift = ft.inverse  # shares temporaries
    one = space_discr.one()
    assert odl.allclose(ift(ft(one)), one)


def test_fourier_trafo_charfun_1d(odl_impl_device_pairs):
    # Characteristic function of [0, 1], its Fourier transform is
    # given by exp(-1j * y / 2) * sinc(y/2)
    impl, device = odl_impl_device_pairs
    def char_interval(x):
        return (x.real >= 0) & (x.real <= 1)

    def char_interval_ft(x):
        if impl == 'numpy':
            return np.exp(-1j * x / 2) * sinc(x / 2) / np.sqrt(2 * np.pi)
        
        elif impl == 'pytorch':
            import torch
            return torch.exp(-1j * x / 2) * sinc_pytorch(x / 2) / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError

    # Base version
    discr = odl.uniform_discr(-2, 2, 40, impl=impl, device=device)
    dft_base = FourierTransform(discr)

    # Complex version, should be as good
    discr = odl.uniform_discr(-2, 2, 40, impl=impl, device=device, dtype='complex64')
    dft_complex = FourierTransform(discr)

    # Without shift
    discr = odl.uniform_discr(-2, 2, 40, impl=impl, device=device, dtype='complex64')
    dft_complex_shift = FourierTransform(discr, shift=False)

    for dft in [dft_base, dft_complex, dft_complex_shift]:
        func_true_ft = dft.range.element(char_interval_ft)
        func_dft = dft(char_interval)
        assert (func_dft - func_true_ft).norm() < 5e-6


def test_fourier_trafo_scaling(odl_impl_device_pairs):
    # Test if the FT scales correctly

    # Characteristic function of [0, 1], its Fourier transform is
    # given by exp(-1j * y / 2) * sinc(y/2)
    impl, device = odl_impl_device_pairs

    def char_interval(x):
        return (x.real >= 0) & (x.real <= 1)

    def char_interval_ft(x):
        if impl == 'numpy':
            return np.exp(-1j * x / 2) * sinc(x / 2) / np.sqrt(2 * np.pi)
        
        elif impl == 'pytorch':
            import torch
            return torch.exp(-1j * x / 2) * sinc_pytorch(x / 2) / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError

    discr = odl.uniform_discr(-2, 2, 40, impl=impl, device=device, dtype='complex128')
    dft = FourierTransform(discr)

    for factor in (2, 1j, -2.5j, 1 - 4j):
        func_true_ft = factor * dft.range.element(char_interval_ft)
        func_dft = dft(factor * discr.element(char_interval))
        assert (func_dft - func_true_ft).norm() < 1e-6


def test_fourier_trafo_sign(impl,odl_real_floating_dtype,odl_impl_device_pairs):
    # Test if the FT sign behaves as expected, i.e. that the FT with sign
    # '+' and '-' have same real parts and opposite imaginary parts.
    impl_backend, device = odl_impl_device_pairs
    discrspace_dtype = complex_dtype(odl_real_floating_dtype)

    def char_interval(x):
        return (x.real >= 0) & (x.real <= 1)

    discr = odl.uniform_discr(-2, 2, 40, impl=impl_backend, device=device, dtype=discrspace_dtype)

    skip_incompatible_impl(impl,discr)

    ft_minus = FourierTransform(discr, sign='-', impl=impl)
    ft_plus = FourierTransform(discr, sign='+', impl=impl)

    func_ft_minus = ft_minus(char_interval)
    func_ft_plus = ft_plus(char_interval)

    data = func_ft_minus.data
    if impl_backend == 'pytorch':
        data = data.detach().cpu().numpy()

    if odl_real_floating_dtype == "float16":
        tolerance = np.linalg.norm(data) * 1e-3
    elif odl_real_floating_dtype == "float32" or odl_real_floating_dtype == float:
        tolerance = np.linalg.norm(data) * 1e-7
    elif odl_real_floating_dtype == "float64" :
        tolerance = np.linalg.norm(data) * 1e-15
    elif odl_real_floating_dtype == "float128":
        if np.__version__<'2':
            # NumPy-1 does not use quadruple precision for the FFT, but double precision
            # and converts the result, so we do not achieve closer tolerance there.
            tolerance = np.linalg.norm(data) * 1e-15
        else:
            tolerance = np.linalg.norm(data) * 1e-19
    else:
        raise TypeError(f"No known tolerance for dtype {odl_real_floating_dtype}")

    def assert_close(x,y):
        x = x.data
        y = y.data
        if impl_backend == 'pytorch':
            x = x.detach().cpu().numpy()
            y = y.detach().cpu().numpy()
        assert(np.linalg.norm(x-y) < tolerance)

    assert_close(func_ft_minus.real, func_ft_plus.real)
    assert_close(func_ft_minus.imag, -func_ft_plus.imag)
    assert_close(ft_minus.inverse.inverse(char_interval), ft_minus(char_interval))
    assert_close(ft_plus.inverse.inverse(char_interval), ft_plus(char_interval))

    discr = odl.uniform_discr(-2, 2, 40, impl=impl_backend, device=device, dtype='float32')
    with pytest.raises(ValueError):
        FourierTransform(discr, sign='+', impl=impl, halfcomplex=True)
    with pytest.raises(ValueError):
        FourierTransform(discr, sign=-1, impl=impl)


def test_fourier_trafo_inverse(impl, sign, odl_impl_device_pairs):
    # Test if the inverse really is the inverse

    impl_backend, device = odl_impl_device_pairs

    def char_interval(x):
        return (x.real >= 0) & (x.real <= 1)

    # Complex-to-complex
    discr = odl.uniform_discr(-2, 2, 40, impl=impl_backend, device=device, dtype='complex64')

    skip_incompatible_impl(impl,discr)

    discr_char = discr.element(char_interval)

    ft = FourierTransform(discr, sign=sign, impl=impl)
    assert all_almost_equal(ft.inverse(ft(char_interval)), discr_char)
    assert all_almost_equal(ft.adjoint(ft(char_interval)), discr_char)

    # Half-complex
    discr = odl.uniform_discr(-2, 2, 40, impl=impl_backend, device=device,dtype='float32')
    ft = FourierTransform(discr, impl=impl, halfcomplex=True)
    assert all_almost_equal(ft.inverse(ft(char_interval)), discr_char.real)

    def char_rect(x):
        return (x[0].real >= 0) & (x[0].real <= 1) & (x[1].real >= 0) & (x[1].real <= 1)

    # 2D with axes, C2C
    discr = odl.uniform_discr([-2, -2], [2, 2], (20, 10), impl=impl_backend,
                              device=device, dtype='complex64')
    discr_rect = discr.element(char_rect)

    for axes in [(0,), 1]:
        ft = FourierTransform(discr, sign=sign, impl=impl, axes=axes)
        assert all_almost_equal(ft.inverse(ft(char_rect)), discr_rect)
        assert all_almost_equal(ft.adjoint(ft(char_rect)), discr_rect)

    # 2D with axes, halfcomplex
    discr = odl.uniform_discr([-2, -2], [2, 2], (20, 10), impl=impl_backend,
                              device=device, dtype='float32')
    discr_rect = discr.element(char_rect)

    for halfcomplex in [False, True]:
        if halfcomplex and sign == '+':
            continue  # cannot mix halfcomplex with sign

        for axes in [(0,), (1,)]:
            ft = FourierTransform(discr, sign=sign, impl=impl, axes=axes,
                                  halfcomplex=halfcomplex)
            assert all_almost_equal(ft.inverse(ft(char_rect)), discr_rect)
            assert all_almost_equal(ft.adjoint(ft(char_rect)), discr_rect)


def test_fourier_trafo_hat_1d(odl_impl_device_pairs):
    # Hat function as used in linear interpolation. It is not so
    # well discretized by nearest neighbor interpolation, so a larger
    # error is to be expected.
    impl, device = odl_impl_device_pairs
    def hat_func(x):
        if impl == 'numpy':
            out = np.where(x < 0, 1 + x, 1 - x)
            out[x < -1] = 0
            out[x > 1] = 0
            return out
        elif impl == 'pytorch':
            import torch
            out = torch.where(x < 0, 1 + x, 1 - x)
            out[x < -1] = 0
            out[x > 1] = 0
            return out
        else:
            raise NotImplementedError

    def hat_func_ft(x):
        if impl == 'numpy':
            return sinc(x / 2) ** 2 / np.sqrt(2 * np.pi)
        elif impl == 'pytorch':
            return sinc_pytorch(x / 2) ** 2 / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError
    # Using a single-precision implementation, should be as good
    discr = odl.uniform_discr(-2, 2, 101, impl=impl, device=device, dtype='float32')
    dft = FourierTransform(discr)
    func_true_ft = dft.range.element(hat_func_ft)
    func_dft = dft(hat_func)
    assert (func_dft - func_true_ft).norm() < 0.001


def test_fourier_trafo_complex_sum(odl_impl_device_pairs):
    # Sum of characteristic function and hat function, both with
    # known FT's.

    impl, device = odl_impl_device_pairs

    def hat_func(x):
        if impl == 'numpy':
            out = np.where(x < 0, 1 + x, 1 - x)
            out[x < -1] = 0
            out[x > 1] = 0
            return out
        elif impl == 'pytorch':
            import torch
            out = torch.where(x.real < 0, 1 + x.real, 1 - x.real)
            out[x.real < -1] = 0
            out[x.real > 1] = 0
            return out
        else:
            raise NotImplementedError

    def hat_func_ft(x):
        if impl == 'numpy':
            return sinc(x / 2) ** 2 / np.sqrt(2 * np.pi)
        elif impl == 'pytorch':
            return sinc_pytorch(x / 2) ** 2 / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError

    def char_interval(x):
        return (x.real >= 0) & (x.real <= 1)

    def char_interval_ft(x):
        if impl == 'numpy':
            return np.exp(-1j * x / 2) * sinc(x / 2) / np.sqrt(2 * np.pi)
        
        elif impl == 'pytorch':
            import torch
            return torch.exp(-1j * x / 2) * sinc_pytorch(x / 2) / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError

    discr = odl.uniform_discr(-2, 2, 200, impl=impl, device=device, dtype='complex128')
    dft = FourierTransform(discr, shift=False)

    func = discr.element(hat_func) + 1j * discr.element(char_interval)
    func_true_ft = (
        dft.range.element(hat_func_ft)
        + 1j * dft.range.element(char_interval_ft)
    )
    func_dft = dft(func)
    assert (func_dft - func_true_ft).norm() < 0.001


def test_fourier_trafo_gaussian_1d(odl_impl_device_pairs):
    # Gaussian function, will be mapped to itself. Truncation error is
    # relatively large, though, we need a large support.

    impl, device = odl_impl_device_pairs

    def gaussian(x):
        if impl == 'numpy':
            return np.exp(-x ** 2 / 2)
        elif impl == 'pytorch':
            import torch
            return torch.exp(-x ** 2 / 2)
        else:
            raise NotImplementedError
    discr = odl.uniform_discr(-10, 10, 201, impl=impl, device=device)
    dft = FourierTransform(discr)
    func_true_ft = dft.range.element(gaussian)
    func_dft = dft(gaussian)
    assert (func_dft - func_true_ft).norm() < 0.001


def test_fourier_trafo_freq_shifted_charfun_1d(odl_impl_device_pairs):
    # Frequency-shifted characteristic function: mult. with
    # exp(-1j * b * x) corresponds to shifting the FT by b.
    impl, device=odl_impl_device_pairs
    
    def fshift_char_interval(x):
        if impl == 'numpy':
            return np.exp(-1j * x * np.pi) * ((x >= -0.5) & (x <= 0.5))
        elif impl== 'pytorch':
             import torch
             return torch.exp(-1j * x * np.pi) * ((x.real >= -0.5) & (x.real <= 0.5))
        else:
            raise NotImplementedError

    def fshift_char_interval_ft(x):
        if impl=='numpy':    
            return sinc((x + np.pi) / 2) / np.sqrt(2 * np.pi)
        elif impl== 'pytorch':
            return sinc_pytorch((x + np.pi) / 2) / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError
    # Number of points is very important here (aliasing)
    discr = odl.uniform_discr(-2, 2, 400, impl=impl, device=device,
                              dtype='complex64')
    dft = FourierTransform(discr)
    func_true_ft = dft.range.element(fshift_char_interval_ft)
    func_dft = dft(fshift_char_interval)
    assert (func_dft - func_true_ft).norm() < 0.001


def test_dft_with_known_pairs_2d(odl_impl_device_pairs):
    impl, device=odl_impl_device_pairs
    # Frequency-shifted product of characteristic functions
    def fshift_char_rect(x):
        # Characteristic function of the cuboid
        # [-1, 1] x [1, 2]
        if impl =='numpy':
            return (x[0] >= -1) & (x[0] <= 1) & (x[1] >= 1) & (x[1] <= 2)
        elif impl == 'pytorch':
            return (x[0].real >= -1) & (x[0].real <= 1) & (x[1].real >= 1) & (x[1].real <= 2)
        else:
            raise NotImplementedError

    def fshift_char_rect_ft(x):
        # FT is a product of shifted and frequency-shifted sinc functions
        # 1st comp.: 2 * sinc(y)
        # 2nd comp.: exp(-1j * y * 3/2) * sinc(y/2)
        # Overall factor: (2 * pi)^(-1)
        if impl =='numpy':
            return (
                2 * sinc(x[0])
                * np.exp(-1j * x[1] * 3 / 2) * sinc(x[1] / 2)
                / (2 * np.pi)
            )
        elif impl == 'pytorch':
            import torch
            return (
                2 * sinc_pytorch(x[0])
                * torch.exp(-1j * x[1] * 3 / 2) * sinc_pytorch(x[1] / 2)
                / (2 * np.pi)
            )

        else:
            raise NotImplementedError

    discr = odl.uniform_discr([-2] * 2, [2] * 2, (100, 400), impl=impl, device=device, dtype='complex64')
    dft = FourierTransform(discr)
    func_true_ft = dft.range.element(fshift_char_rect_ft)
    func_dft = dft(fshift_char_rect)
    assert (func_dft - func_true_ft).norm() < 0.001


def test_fourier_trafo_completely(odl_impl_device_pairs):
    # Complete explicit test of all FT components on two small examples
    impl, device = odl_impl_device_pairs
    
    # Discretization with 4 points
    discr = odl.uniform_discr(-2, 2, 4, impl=impl, device=device, dtype=complex)
    # Interval boundaries -2, -1, 0, 1, 2
    assert np.allclose(discr.partition.cell_boundary_vecs[0],
                       [-2, -1, 0, 1, 2])
    # Grid points -1.5, -0.5, 0.5, 1.5
    assert np.allclose(discr.grid.coord_vectors[0],
                       [-1.5, -0.5, 0.5, 1.5])

    # First test function, symmetric. Can be represented exactly in the
    # discretization.
    def f(x):
        if impl == 'numpy':
            return (x >= -1) & (x <= 1)
        elif impl == 'pytorch':
            return (x.real >= -1) & (x.real <= 1)
        else:
            raise NotImplementedError


    def fhat(x):
        if impl == 'numpy':
            return np.sqrt(2 / np.pi) * sinc(x)
        elif impl == 'pytorch':
            return  np.sqrt(2 / np.pi) * sinc_pytorch(x)
        else:
            raise NotImplementedError

    # Discretize f, check values
    f_discr = discr.element(f)
    assert all_almost_equal(f_discr, [0, 1, 1, 0])

    # "s" = shifted, "n" = not shifted

    # Reciprocal grids
    recip_s = reciprocal_grid(discr.grid, shift=True)
    recip_n = reciprocal_grid(discr.grid, shift=False)
    assert np.allclose(recip_s.coord_vectors[0],
                       np.linspace(-np.pi, np.pi / 2, 4))
    assert np.allclose(recip_n.coord_vectors[0],
                       np.linspace(-3 * np.pi / 4, 3 * np.pi / 4, 4))

    # Range
    range_part_s = odl.uniform_partition_fromgrid(recip_s)
    range_s = odl.uniform_discr_frompartition(range_part_s, dtype=complex, impl=impl, device=device)
    range_part_n = odl.uniform_partition_fromgrid(recip_n)
    range_n = odl.uniform_discr_frompartition(range_part_n, dtype=complex, impl=impl, device=device)

    namespace = discr.array_namespace
    backend = discr.array_backend

    # Pre-processing
    preproc_s = [1, -1, 1, -1]
    preproc_n = [np.exp(1j * 3 / 4 * np.pi * k) for k in range(4)]

    fpre_s = dft_preprocess_data(f_discr, shift=True)
    fpre_n = dft_preprocess_data(f_discr, shift=False)
    assert all_almost_equal(fpre_s, f_discr * discr.element(preproc_s))
    assert all_almost_equal(fpre_n, f_discr * discr.element(preproc_n))

    # FFT step, replicating the _call_numpy method
    fft_s = namespace.fft.fftn(fpre_s, s=discr.shape, axes=[0])
    fft_n = namespace.fft.fftn(fpre_n, s=discr.shape, axes=[0])
    s = backend.array_constructor([0, -1 + 1j, 2, -1 - 1j], device=device, dtype=complex)
    n = backend.array_constructor([
            np.exp(1j * np.pi * (3 - 2 * k) / 4)
            + np.exp(1j * np.pi * (3 - 2 * k) / 2)
            for k in range(4)], device=device, dtype=complex
         )
    assert namespace.allclose(fft_s, s)
    assert namespace.allclose(fft_n, n)

    # Interpolation kernel FT
    interp_s = np.sinc(np.linspace(-1 / 2, 1 / 4, 4)) / np.sqrt(2 * np.pi)
    interp_n = np.sinc(np.linspace(-3 / 8, 3 / 8, 4)) / np.sqrt(2 * np.pi)
    assert np.allclose(interp_s,
                       _interp_kernel_ft(np.linspace(-1 / 2, 1 / 4, 4),
                                         interp='nearest'))
    assert np.allclose(interp_n,
                       _interp_kernel_ft(np.linspace(-3 / 8, 3 / 8, 4),
                                         interp='nearest'))

    # Post-processing
    postproc_s = np.exp(1j * np.pi * np.linspace(-3 / 2, 3 / 4, 4))
    postproc_n = np.exp(1j * np.pi * np.linspace(-9 / 8, 9 / 8, 4))

    fpost_s = dft_postprocess_data(
        range_s.element(fft_s), real_grid=discr.grid, recip_grid=recip_s,
        shift=[True], axes=(0,), interp='nearest')
    fpost_n = dft_postprocess_data(
        range_n.element(fft_n), real_grid=discr.grid, recip_grid=recip_n,
        shift=[False], axes=(0,), interp='nearest')
    
    postproc_s = backend.array_constructor(postproc_s, device=device)
    postproc_n = backend.array_constructor(postproc_n, device=device)

    interp_s = backend.array_constructor(interp_s, device=device)
    interp_n = backend.array_constructor(interp_n, device=device)

    assert namespace.allclose(fpost_s, fft_s * postproc_s * interp_s)
    assert namespace.allclose(fpost_n, fft_n * postproc_n * interp_n)

    # Comparing to the known result sqrt(2/pi) * sinc(x)
    recip_s_array = backend.array_constructor(recip_s.coord_vectors[0], device=device)
    recip_n_array = backend.array_constructor(recip_n.coord_vectors[0], device=device)
    assert namespace.allclose(fpost_s.real, fhat(recip_s_array))
    assert namespace.allclose(fpost_n.real, fhat(recip_n_array))

    # Doing the exact same with direct application of the FT operator
    ft_op_s = FourierTransform(discr, shift=True)
    ft_op_n = FourierTransform(discr, shift=False)
    assert ft_op_s.range.grid == recip_s
    assert ft_op_n.range.grid == recip_n

    ft_f_s = ft_op_s(f)
    ft_f_n = ft_op_n(f)
    assert all_almost_equal(ft_f_s.real, fhat(recip_s_array))
    assert all_almost_equal(ft_f_n.real, fhat(recip_n_array))

    # Second test function, asymmetric. Can also be represented exactly in the
    # discretization.
    def f(x):
        return (x.real >= 0) & (x.real <= 1)

    def fhat(x):
        if impl == 'numpy':
            return np.exp(-1j * x / 2) * sinc(x / 2) / np.sqrt(2 * np.pi)
        elif impl == 'pytorch':
            return namespace.exp(-1j * x / 2) * sinc_pytorch(x / 2) / np.sqrt(2 * np.pi)
        else:
            raise NotImplementedError
    # Discretize f, check values
    f_discr = discr.element(f)
    assert all_almost_equal(f_discr, [0, 0, 1, 0])

    # Pre-processing
    fpre_s = dft_preprocess_data(f_discr, shift=True)
    fpre_n = dft_preprocess_data(f_discr, shift=False)
    assert all_almost_equal(fpre_s, [0, 0, 1, 0])
    assert all_almost_equal(fpre_n, backend.array_constructor([0, 0, -1j, 0], device=device, dtype=complex))

    # FFT step
    fft_s = namespace.fft.fftn(fpre_s, s=discr.shape, axes=[0])
    fft_n = namespace.fft.fftn(fpre_n, s=discr.shape, axes=[0])
    assert all_almost_equal(fft_s, [1, -1, 1, -1])
    assert all_almost_equal(fft_n, backend.array_constructor([-1j, 1j, -1j, 1j], device=device, dtype=complex))

    fpost_s = dft_postprocess_data(
        range_s.element(fft_s), real_grid=discr.grid, recip_grid=recip_s,
        shift=[True], axes=(0,), interp='nearest')
    fpost_n = dft_postprocess_data(
        range_n.element(fft_n), real_grid=discr.grid, recip_grid=recip_n,
        shift=[False], axes=(0,), interp='nearest')

    assert all_almost_equal(fpost_s, fft_s * postproc_s * interp_s)
    assert all_almost_equal(fpost_n, fft_n * postproc_n * interp_n)

    # Comparing to the known result exp(-1j*x/2) * sinc(x/2) / sqrt(2*pi)
    recip_s = backend.array_constructor(recip_s.coord_vectors[0], device=device)
    recip_n = backend.array_constructor(recip_n.coord_vectors[0], device=device)
    assert all_almost_equal(fpost_s, fhat(recip_s))
    assert all_almost_equal(fpost_n, fhat(recip_n))

    # Doing the exact same with direct application of the FT operator
    ft_f_s = ft_op_s(f)
    ft_f_n = ft_op_n(f)
    assert all_almost_equal(ft_f_s, fhat(recip_s))
    assert all_almost_equal(ft_f_n, fhat(recip_n))


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