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)

[((2105, 2147), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""width"""', '[1, 2.3]'], {}), "('width', [1, 2.3])\n", (2128, 2147), False, 'import pytest\n'), ((2941, 2987), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""width"""', "[-1, 0, 'a']"], {}), "('width', [-1, 0, 'a'])\n", (2964, 2987), False, 'import pytest\n'), ((3343, 3386), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""stddev"""', '[1, 2.3]'], {}), "('stddev', [1, 2.3])\n", (3366, 3386), False, 'import pytest\n'), ((4227, 4274), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""stddev"""', "[-1, 0, 'a']"], {}), "('stddev', [-1, 0, 'a'])\n", (4250, 4274), False, 'import pytest\n'), ((4661, 4704), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""stddev"""', '[1, 2.3]'], {}), "('stddev', [1, 2.3])\n", (4684, 4704), False, 'import pytest\n'), ((5560, 5607), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""stddev"""', "[-1, 0, 'a']"], {}), "('stddev', [-1, 0, 'a'])\n", (5583, 5607), False, 'import pytest\n'), ((5990, 6033), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""width"""', '[1, 3, 9]'], {}), "('width', [1, 3, 9])\n", (6013, 6033), False, 'import pytest\n'), ((6780, 6826), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""width"""', "[-1, 0, 'a']"], {}), "('width', [-1, 0, 'a'])\n", (6803, 6826), False, 'import pytest\n'), ((1191, 1230), 'numpy.allclose', 'np.allclose', (['flux_smooth1', 'flux_smooth2'], {}), '(flux_smooth1, flux_smooth2)\n', (1202, 1230), True, 'import numpy as np\n'), ((1689, 1720), 'numpy.array', 'np.array', (['[0.5, 1, 2, 0.5, 0.2]'], {}), '([0.5, 1, 2, 0.5, 0.2])\n', (1697, 1720), True, 'import numpy as np\n'), ((1797, 1835), 'astropy.convolution.CustomKernel', 'convolution.CustomKernel', (['numpy_kernel'], {}), '(numpy_kernel)\n', (1821, 1835), False, 'from astropy import convolution\n'), ((1864, 1914), 'astropy.convolution.convolve', 'convolution.convolve', (['flux_original', 'custom_kernel'], {}), '(flux_original, custom_kernel)\n', (1884, 1914), False, 'from astropy import convolution\n'), ((2502, 2532), 'astropy.convolution.Box1DKernel', 'convolution.Box1DKernel', (['width'], {}), '(width)\n', (2525, 2532), False, 'from astropy import convolution\n'), ((2561, 2608), 'astropy.convolution.convolve', 'convolution.convolve', (['flux_original', 'box_kernel'], {}), '(flux_original, box_kernel)\n', (2581, 2608), False, 'from astropy import convolution\n'), ((3762, 3798), 'astropy.convolution.Gaussian1DKernel', 'convolution.Gaussian1DKernel', (['stddev'], {}), '(stddev)\n', (3790, 3798), False, 'from astropy import convolution\n'), ((3827, 3879), 'astropy.convolution.convolve', 'convolution.convolve', (['flux_original', 'gaussian_kernel'], {}), '(flux_original, gaussian_kernel)\n', (3847, 3879), False, 'from astropy import convolution\n'), ((5102, 5139), 'astropy.convolution.Trapezoid1DKernel', 'convolution.Trapezoid1DKernel', (['stddev'], {}), '(stddev)\n', (5131, 5139), False, 'from astropy import convolution\n'), ((5168, 5221), 'astropy.convolution.convolve', 'convolution.convolve', (['flux_original', 'trapezoid_kernel'], {}), '(flux_original, trapezoid_kernel)\n', (5188, 5221), False, 'from astropy import convolution\n'), ((6408, 6437), 'scipy.signal.medfilt', 'medfilt', (['flux_original', 'width'], {}), '(flux_original, width)\n', (6415, 6437), False, 'from scipy.signal import medfilt\n'), ((1755, 1775), 'numpy.sum', 'np.sum', (['numpy_kernel'], {}), '(numpy_kernel)\n', (1761, 1775), True, 'import numpy as np\n'), ((3280, 3305), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (3293, 3305), False, 'import pytest\n'), ((4592, 4617), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (4605, 4617), False, 'import pytest\n'), ((5920, 5945), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (5933, 5945), False, 'import pytest\n'), ((7111, 7136), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (7124, 7136), False, 'import pytest\n')]