import multiprocessing
import pytest
import numpy as np
from numpy.testing import assert_array_almost_equal

from pylops.utils import dottest
from pylops.signalprocessing import Radon2D, Radon3D
from pylops.optimization.sparsity import FISTA

par1 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx': 21, 'npy': 17,
        'pymax': 1e-2, 'pxmax': 2e-2,
        'centeredh': True, 'kind': 'linear', 'interp': True,
        'engine': 'numpy'}  # linear, centered, linear interp, numpy
par2 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx': 21, 'npy': 17,
        'pymax': 1e-2, 'pxmax': 2e-2,
        'centeredh': False, 'kind': 'linear', 'interp': True,
        'engine': 'numpy'}  # linear, uncentered, linear interp, numpy
par3 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx':21, 'npy':17,
        'pymax':1e-2, 'pxmax':2e-2,
        'centeredh': True, 'kind': 'linear', 'interp':True,
        'engine':'numba'} # linear, centered, linear interp, numba
par4 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx':21, 'npy':17,
        'pymax':1e-2, 'pxmax':2e-2,
        'centeredh': False, 'kind': 'linear', 'interp':False,
        'engine':'numba'}  # linear, uncentered, linear interp, numba
par5 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx':21, 'npy':17,
        'pymax': 8e-3, 'pxmax': 7e-3,
        'centeredh': True, 'kind': 'parabolic', 'interp':False,
        'engine':'numpy'}  # parabolic, centered, no interp, numpy
par6 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx':21, 'npy':17,
        'pymax': 8e-3, 'pxmax': 7e-3,
        'centeredh': False, 'kind': 'parabolic', 'interp':True,
        'engine':'numba'}  # parabolic, uncentered, interp, numba
par7 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx':21, 'npy':17,
        'pymax': 9e-2, 'pxmax': 8e-2,
        'centeredh': True, 'kind': 'hyperbolic', 'interp':True,
        'engine':'numpy'}  # hyperbolic, centered, interp, numpy
par8 = {'nt': 11, 'nhx': 21, 'nhy': 10, 'npx':21, 'npy':17,
        'pymax': 7e-2, 'pxmax': 8e-2,
        'centeredh': False, 'kind': 'hyperbolic', 'interp':False,
        'engine':'numba'}  # hyperbolic, uncentered, interp, numba


def test_unknown_engine():
    """Check error is raised if unknown engine is passed
    """
    with pytest.raises(KeyError):
        _ = Radon2D(None, None, None, engine='foo')
    with pytest.raises(KeyError):
        _ = Radon3D(None, None, None, None, None, engine='foo')


@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4),
                                 (par5), (par6), (par7), (par8)])
def test_Radon2D(par):
    """Dot-test, forward and adjoint consistency check
    (for onthefly parameter), and sparse inverse for Radon2D operator
    """
    dt, dh = 0.005, 1
    t = np.arange(par['nt']) * dt
    h = np.arange(par['nhx']) * dh
    px = np.linspace(0, par['pxmax'], par['npx'])
    x = np.zeros((par['npx'], par['nt']))
    x[2, par['nt']//2] = 1

    Rop = Radon2D(t, h, px, centeredh=par['centeredh'],
                  interp=par['interp'], kind=par['kind'],
                  onthefly=False, engine=par['engine'],
                  dtype='float64')
    R1op = Radon2D(t, h, px, centeredh=par['centeredh'],
                   interp=par['interp'], kind=par['kind'],
                   onthefly=True, engine=par['engine'],
                   dtype='float64')
    assert dottest(Rop, par['nhx']*par['nt'], par['npx']*par['nt'], tol=1e-3)

    y = Rop * x.ravel()
    y1 = R1op * x.ravel()
    assert_array_almost_equal(y, y1, decimal=4)

    xadj = Rop.H * y
    xadj1 = R1op.H * y
    assert_array_almost_equal(xadj, xadj1, decimal=4)

    xinv, _, _ = FISTA(Rop, y, 30, eps=1e0, returninfo=True)
    assert_array_almost_equal(x.ravel(), xinv, decimal=1)


@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4),
                                 (par5), (par6), (par7), (par8)])
def test_Radon3D(par):
    """Dot-test,  forward and adjoint consistency check
    (for onthefly parameter), and sparse inverse for Radon3D operator
    """
    if par['engine'] == 'numpy' or \
        multiprocessing.cpu_count() >= 4: # avoid timeout in travis for numba

        dt, dhy, dhx = 0.005, 1, 1
        t = np.arange(par['nt']) * dt
        hy = np.arange(par['nhy']) * dhy
        hx = np.arange(par['nhx']) * dhx
        py = np.linspace(0, par['pymax'], par['npy'])
        px = np.linspace(0, par['pxmax'], par['npx'])
        x = np.zeros((par['npy'], par['npx'], par['nt']))
        x[3, 2, par['nt']//2] = 1

        Rop = Radon3D(t, hy, hx, py, px, centeredh=par['centeredh'],
                      interp=par['interp'], kind=par['kind'],
                      onthefly=False, engine=par['engine'],
                      dtype='float64')
        R1op = Radon3D(t, hy, hx, py, px, centeredh=par['centeredh'],
                       interp=par['interp'], kind=par['kind'],
                       onthefly=True, engine=par['engine'],
                       dtype='float64')

        assert dottest(Rop, par['nhy']*par['nhx']*par['nt'],
                       par['npy']*par['npx']*par['nt'], tol=1e-3)
        y = Rop * x.flatten()
        y1 = R1op * x.flatten()
        assert_array_almost_equal(y, y1, decimal=4)

        xadj = Rop.H * y
        xadj1 = R1op.H * y
        assert_array_almost_equal(xadj, xadj1, decimal=4)

        if Rop.engine == 'numba': # as numpy is too slow here...
            xinv, _, _ = FISTA(Rop, y, 200, eps=3e0, returninfo=True)
            assert_array_almost_equal(x.flatten(), xinv, decimal=1)
