import numpy as np
import pytest
from astropy import convolution
from scipy.signal import medfilt
import astropy.units as u
from ..spectra.spectrum1d import Spectrum1D
from ..tests.spectral_examples import simulated_spectra

from ..manipulation.smoothing import (convolution_smooth, box_smooth,
                                      gaussian_smooth, trapezoid_smooth,
                                      median_smooth)

def compare_flux(flux_smooth1, flux_smooth2, flux_original, rtol=0.01):
    """
    There are two things to compare for each set of smoothing:

    1. Compare the smoothed flux from the astropy machinery vs
       the smoothed flux from specutils.  This is done by
       comparing flux_smooth1 and flux_smooth2.

    2. Next we want to compare the smoothed flux to the original
       flux.  This is a little more difficult as smoothing will
       make a difference for median filter, but less so for
       convolution based smoothing if the kernel is normalized
       (area under the kernel = 1).

       In this second case the rtol (relative tolerance) is used
       judiciously.

    """

    # Compare, element by element, the two smoothed fluxes.
    assert np.allclose(flux_smooth1, flux_smooth2)

    # Compare the total spectral flux of the smoothed to the original.
    assert np.allclose(sum(flux_smooth1), sum(flux_original), rtol=rtol)


def test_smooth_custom_kernel(simulated_spectra):
    """
    Test CustomKernel smoothing with correct parmaeters.
    """

    # Create the original spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1
    flux_original = spec1.flux

    # Create a custom kernel (some weird asymmetric-ness)
    numpy_kernel = np.array([0.5, 1, 2, 0.5, 0.2])
    numpy_kernel = numpy_kernel / np.sum(numpy_kernel)

    custom_kernel = convolution.CustomKernel(numpy_kernel)
    flux_smoothed_astropy = convolution.convolve(flux_original, custom_kernel)

    # Calculate the custom smoothed
    spec1_smoothed = convolution_smooth(spec1, custom_kernel)
    compare_flux(spec1_smoothed.flux.value, flux_smoothed_astropy, flux_original.value)


@pytest.mark.parametrize("width", [1, 2.3])
def test_smooth_box_good(simulated_spectra, width):
    """
    Test Box1DKernel smoothing with correct parmaeters.

    Width values need to be a number greater than 0.
    """

    # Create the original spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1
    flux_original = spec1.flux

    # Calculate the smoothed flux using Astropy
    box_kernel = convolution.Box1DKernel(width)
    flux_smoothed_astropy = convolution.convolve(flux_original, box_kernel)

    # Calculate the box smoothed
    spec1_smoothed = box_smooth(spec1, width)
    compare_flux(spec1_smoothed.flux.value, flux_smoothed_astropy, flux_original.value)

    # Check the input and output units
    assert spec1.wavelength.unit == spec1_smoothed.wavelength.unit
    assert spec1.flux.unit == spec1_smoothed.flux.unit

@pytest.mark.parametrize("width", [-1, 0, 'a'])
def test_smooth_box_bad(simulated_spectra, width):
    """
    Test Box1DKernel smoothing with incorrect parmaeters.

    Width values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1

    # Test bad input parameters
    with pytest.raises(ValueError):
        box_smooth(spec1, width)


@pytest.mark.parametrize("stddev", [1, 2.3])
def test_smooth_gaussian_good(simulated_spectra, stddev):
    """
    Test Gaussian1DKernel smoothing with correct parmaeters.

    Standard deviation values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1
    flux_original = spec1.flux

    # Calculate the smoothed flux using Astropy
    gaussian_kernel = convolution.Gaussian1DKernel(stddev)
    flux_smoothed_astropy = convolution.convolve(flux_original, gaussian_kernel)

    # Test gaussian smoothing
    spec1_smoothed = gaussian_smooth(spec1, stddev)
    compare_flux(spec1_smoothed.flux.value, flux_smoothed_astropy, flux_original.value, rtol=0.02)

    # Check the input and output units
    assert spec1.wavelength.unit == spec1_smoothed.wavelength.unit
    assert spec1.flux.unit == spec1_smoothed.flux.unit


@pytest.mark.parametrize("stddev", [-1, 0, 'a'])
def test_smooth_gaussian_bad(simulated_spectra, stddev):
    """
    Test MexicanHat1DKernel smoothing with incorrect parmaeters.

    Standard deviation values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1

    # Test bad input paramters
    with pytest.raises(ValueError):
        gaussian_smooth(spec1, stddev)


@pytest.mark.parametrize("stddev", [1, 2.3])
def test_smooth_trapezoid_good(simulated_spectra, stddev):
    """
    Test Trapezoid1DKernel smoothing with correct parmaeters.

    Standard deviation values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1
    flux_original = spec1.flux

    # Create the flux_smoothed which is what we want to compare to
    trapezoid_kernel = convolution.Trapezoid1DKernel(stddev)
    flux_smoothed_astropy = convolution.convolve(flux_original, trapezoid_kernel)

    # Test trapezoid smoothing
    spec1_smoothed = trapezoid_smooth(spec1, stddev)
    compare_flux(spec1_smoothed.flux.value, flux_smoothed_astropy, flux_original.value)

    # Check the input and output units
    assert spec1.wavelength.unit == spec1_smoothed.wavelength.unit
    assert spec1.flux.unit == spec1_smoothed.flux.unit


@pytest.mark.parametrize("stddev", [-1, 0, 'a'])
def test_smooth_trapezoid_bad(simulated_spectra, stddev):
    """
    Test Trapezoid1DKernel smoothing with incorrect parmaeters.

    Standard deviation values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1

    # Test bad parameters
    with pytest.raises(ValueError):
        trapezoid_smooth(spec1, stddev)


@pytest.mark.parametrize("width", [1, 3, 9])
def test_smooth_median_good(simulated_spectra, width):
    """
    Test Median smoothing with correct parmaeters.

    Width values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1
    flux_original = spec1.flux

    # Create the flux_smoothed which is what we want to compare to
    flux_smoothed_astropy = medfilt(flux_original, width)

    # Test median smoothing
    spec1_smoothed = median_smooth(spec1, width)
    compare_flux(spec1_smoothed.flux.value, flux_smoothed_astropy, flux_original.value, rtol=0.15)

    # Check the input and output units
    assert spec1.wavelength.unit == spec1_smoothed.wavelength.unit
    assert spec1.flux.unit == spec1_smoothed.flux.unit


@pytest.mark.parametrize("width", [-1, 0, 'a'])
def test_smooth_median_bad(simulated_spectra, width):
    """
    Test Median smoothing with incorrect parmaeters.

    Width values need to be a number greater than 0.
    """

    #  Create the spectrum
    spec1 = simulated_spectra.s1_um_mJy_e1

    # Test bad parameters
    with pytest.raises(ValueError):
        median_smooth(spec1, width)

[((60, 1, 60, 43), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(60, 25, 60, 32): '"""width"""', (60, 34, 60, 42): '[1, 2.3]'}, {}), "('width', [1, 2.3])", False, 'import pytest\n'), ((84, 1, 84, 47), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(84, 25, 84, 32): '"""width"""', (84, 34, 84, 46): "[-1, 0, 'a']"}, {}), "('width', [-1, 0, 'a'])", False, 'import pytest\n'), ((100, 1, 100, 44), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(100, 25, 100, 33): '"""stddev"""', (100, 35, 100, 43): '[1, 2.3]'}, {}), "('stddev', [1, 2.3])", False, 'import pytest\n'), ((125, 1, 125, 48), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(125, 25, 125, 33): '"""stddev"""', (125, 35, 125, 47): "[-1, 0, 'a']"}, {}), "('stddev', [-1, 0, 'a'])", False, 'import pytest\n'), ((141, 1, 141, 44), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(141, 25, 141, 33): '"""stddev"""', (141, 35, 141, 43): '[1, 2.3]'}, {}), "('stddev', [1, 2.3])", False, 'import pytest\n'), ((166, 1, 166, 48), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(166, 25, 166, 33): '"""stddev"""', (166, 35, 166, 47): "[-1, 0, 'a']"}, {}), "('stddev', [-1, 0, 'a'])", False, 'import pytest\n'), ((182, 1, 182, 44), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(182, 25, 182, 32): '"""width"""', (182, 34, 182, 43): '[1, 3, 9]'}, {}), "('width', [1, 3, 9])", False, 'import pytest\n'), ((206, 1, 206, 47), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(206, 25, 206, 32): '"""width"""', (206, 34, 206, 46): "[-1, 0, 'a']"}, {}), "('width', [-1, 0, 'a'])", False, 'import pytest\n'), ((33, 11, 33, 50), 'numpy.allclose', 'np.allclose', ({(33, 23, 33, 35): 'flux_smooth1', (33, 37, 33, 49): 'flux_smooth2'}, {}), '(flux_smooth1, flux_smooth2)', True, 'import numpy as np\n'), ((49, 19, 49, 50), 'numpy.array', 'np.array', ({(49, 28, 49, 49): '[0.5, 1, 2, 0.5, 0.2]'}, {}), '([0.5, 1, 2, 0.5, 0.2])', True, 'import numpy as np\n'), ((52, 20, 52, 58), 'astropy.convolution.CustomKernel', 'convolution.CustomKernel', ({(52, 45, 52, 57): 'numpy_kernel'}, {}), '(numpy_kernel)', False, 'from astropy import convolution\n'), ((53, 28, 53, 78), 'astropy.convolution.convolve', 'convolution.convolve', ({(53, 49, 53, 62): 'flux_original', (53, 64, 53, 77): 'custom_kernel'}, {}), '(flux_original, custom_kernel)', False, 'from astropy import convolution\n'), ((73, 17, 73, 47), 'astropy.convolution.Box1DKernel', 'convolution.Box1DKernel', ({(73, 41, 73, 46): 'width'}, {}), '(width)', False, 'from astropy import convolution\n'), ((74, 28, 74, 75), 'astropy.convolution.convolve', 'convolution.convolve', ({(74, 49, 74, 62): 'flux_original', (74, 64, 74, 74): 'box_kernel'}, {}), '(flux_original, box_kernel)', False, 'from astropy import convolution\n'), ((113, 22, 113, 58), 'astropy.convolution.Gaussian1DKernel', 'convolution.Gaussian1DKernel', ({(113, 51, 113, 57): 'stddev'}, {}), '(stddev)', False, 'from astropy import convolution\n'), ((114, 28, 114, 80), 'astropy.convolution.convolve', 'convolution.convolve', ({(114, 49, 114, 62): 'flux_original', (114, 64, 114, 79): 'gaussian_kernel'}, {}), '(flux_original, gaussian_kernel)', False, 'from astropy import convolution\n'), ((154, 23, 154, 60), 'astropy.convolution.Trapezoid1DKernel', 'convolution.Trapezoid1DKernel', ({(154, 53, 154, 59): 'stddev'}, {}), '(stddev)', False, 'from astropy import convolution\n'), ((155, 28, 155, 81), 'astropy.convolution.convolve', 'convolution.convolve', ({(155, 49, 155, 62): 'flux_original', (155, 64, 155, 80): 'trapezoid_kernel'}, {}), '(flux_original, trapezoid_kernel)', False, 'from astropy import convolution\n'), ((195, 28, 195, 57), 'scipy.signal.medfilt', 'medfilt', ({(195, 36, 195, 49): 'flux_original', (195, 51, 195, 56): 'width'}, {}), '(flux_original, width)', False, 'from scipy.signal import medfilt\n'), ((50, 34, 50, 54), 'numpy.sum', 'np.sum', ({(50, 41, 50, 53): 'numpy_kernel'}, {}), '(numpy_kernel)', True, 'import numpy as np\n'), ((96, 9, 96, 34), 'pytest.raises', 'pytest.raises', ({(96, 23, 96, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n'), ((137, 9, 137, 34), 'pytest.raises', 'pytest.raises', ({(137, 23, 137, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n'), ((178, 9, 178, 34), 'pytest.raises', 'pytest.raises', ({(178, 23, 178, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n'), ((218, 9, 218, 34), 'pytest.raises', 'pytest.raises', ({(218, 23, 218, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n')]