# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Tests for the core module.
"""

import itertools
import warnings

from astropy.modeling.models import Gaussian1D, Gaussian2D
from astropy.utils.exceptions import (AstropyDeprecationWarning,
                                      AstropyUserWarning)
import numpy as np
from numpy.testing import assert_allclose
import pytest

from ..gaussian import (centroid_1dg, centroid_2dg, gaussian1d_moments,
                        fit_2dgaussian)

try:
    # the fitting routines in astropy use scipy.optimize
    import scipy  # noqa
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False


XCEN = 25.7
YCEN = 26.2
XSTDS = [3.2, 4.0]
YSTDS = [5.7, 4.1]
THETAS = np.array([30., 45.]) * np.pi / 180.

DATA = np.zeros((3, 3))
DATA[0:2, 1] = 1.
DATA[1, 0:2] = 1.
DATA[1, 1] = 2.


@pytest.mark.skipif('not HAS_SCIPY')
@pytest.mark.parametrize(('x_std', 'y_std', 'theta'),
                         list(itertools.product(XSTDS, YSTDS, THETAS)))
def test_centroids(x_std, y_std, theta):
    model = Gaussian2D(2.4, XCEN, YCEN, x_stddev=x_std, y_stddev=y_std,
                       theta=theta)
    y, x = np.mgrid[0:50, 0:47]
    data = model(x, y)
    xc, yc = centroid_1dg(data)
    assert_allclose((xc, yc), (XCEN, YCEN), rtol=0, atol=1.e-3)
    xc, yc = centroid_2dg(data)
    assert_allclose((xc, yc), (XCEN, YCEN), rtol=0, atol=1.e-3)

    # test with errors
    error = np.sqrt(data)
    xc, yc = centroid_1dg(data, error=error)
    assert_allclose((xc, yc), (XCEN, YCEN), rtol=0, atol=1.e-3)
    xc, yc = centroid_2dg(data, error=error)
    assert_allclose((xc, yc), (XCEN, YCEN), rtol=0, atol=1.e-3)

    # test with mask
    mask = np.zeros(data.shape, dtype=bool)
    data[10, 10] = 1.e5
    mask[10, 10] = True
    xc, yc = centroid_1dg(data, mask=mask)
    assert_allclose((xc, yc), (XCEN, YCEN), rtol=0, atol=1.e-3)
    xc, yc = centroid_2dg(data, mask=mask)
    assert_allclose((xc, yc), (XCEN, YCEN), rtol=0, atol=1.e-3)


@pytest.mark.skipif('not HAS_SCIPY')
@pytest.mark.parametrize('use_mask', [True, False])
def test_centroids_nan_withmask(use_mask):
    xc_ref = 24.7
    yc_ref = 25.2
    model = Gaussian2D(2.4, xc_ref, yc_ref, x_stddev=5.0, y_stddev=5.0)
    y, x = np.mgrid[0:50, 0:50]
    data = model(x, y)
    data[20, :] = np.nan
    if use_mask:
        mask = np.zeros(data.shape, dtype=bool)
        mask[20, :] = True
        nwarn = 0
    else:
        mask = None
        nwarn = 1

    with warnings.catch_warnings(record=True) as warnlist:
        xc, yc = centroid_1dg(data, mask=mask)
        assert_allclose([xc, yc], [xc_ref, yc_ref], rtol=0, atol=1.e-3)
    assert len(warnlist) == nwarn
    if nwarn == 1:
        assert issubclass(warnlist[0].category, AstropyUserWarning)

    with warnings.catch_warnings(record=True) as warnlist:
        xc, yc = centroid_2dg(data, mask=mask)
        assert_allclose([xc, yc], [xc_ref, yc_ref], rtol=0, atol=1.e-3)
    assert len(warnlist) == nwarn
    if nwarn == 1:
        assert issubclass(warnlist[0].category, AstropyUserWarning)


@pytest.mark.skipif('not HAS_SCIPY')
def test_invalid_mask_shape():
    data = np.zeros((4, 4))
    mask = np.zeros((2, 2), dtype=bool)

    with pytest.raises(ValueError):
        centroid_1dg(data, mask=mask)
    with pytest.raises(ValueError):
        centroid_2dg(data, mask=mask)
    with pytest.raises(ValueError):
        gaussian1d_moments(data, mask=mask)


@pytest.mark.skipif('not HAS_SCIPY')
def test_invalid_error_shape():
    error = np.zeros((2, 2), dtype=bool)
    with pytest.raises(ValueError):
        centroid_1dg(np.zeros((4, 4)), error=error)
    with pytest.raises(ValueError):
        centroid_2dg(np.zeros((4, 4)), error=error)


@pytest.mark.skipif('not HAS_SCIPY')
def test_centroid_2dg_dof():
    data = np.ones((2, 2))
    with pytest.raises(ValueError):
        centroid_2dg(data)


def test_gaussian1d_moments():
    x = np.arange(100)
    desired = (75, 50, 5)
    g = Gaussian1D(*desired)
    data = g(x)
    result = gaussian1d_moments(data)
    assert_allclose(result, desired, rtol=0, atol=1.e-6)

    data[0] = 1.e5
    mask = np.zeros(data.shape).astype(bool)
    mask[0] = True
    result = gaussian1d_moments(data, mask=mask)
    assert_allclose(result, desired, rtol=0, atol=1.e-6)

    data[0] = np.nan
    mask = np.zeros(data.shape).astype(bool)
    mask[0] = True
    with warnings.catch_warnings(record=True) as warnlist:
        result = gaussian1d_moments(data, mask=mask)
        assert_allclose(result, desired, rtol=0, atol=1.e-6)
    assert len(warnlist) == 1
    assert issubclass(warnlist[0].category, AstropyUserWarning)


@pytest.mark.skipif('not HAS_SCIPY')
def test_fitgaussian2d():
    model = Gaussian2D(2.4, XCEN, YCEN, x_stddev=3.2, y_stddev=5.7,
                       theta=np.pi/6.)
    y, x = np.mgrid[0:50, 0:47]
    data = model(x, y)
    error = np.ones(data.shape, dtype=float)
    mask = np.zeros(data.shape, dtype=bool)
    with pytest.warns(AstropyDeprecationWarning):
        gfit = fit_2dgaussian(data, error=error, mask=mask)
    assert_allclose(gfit.x_mean.value, XCEN)
    assert_allclose(gfit.y_mean.value, YCEN)
