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

import astropy.units as u
import numpy as np
import pytest
from astropy.modeling.models import Gaussian1D, Gaussian2D
from astropy.table import QTable
from astropy.utils.exceptions import AstropyUserWarning
from numpy.testing import assert_allclose

from photutils.psf import CircularGaussianPRF, make_psf_model_image
from photutils.psf.utils import (_get_psf_model_main_params,
                                 _interpolate_missing_data,
                                 _validate_psf_model, fit_2dgaussian, fit_fwhm)


@pytest.fixture(name='test_data')
def fixture_test_data():
    psf_model = CircularGaussianPRF()
    model_shape = (9, 9)
    n_sources = 10
    shape = (101, 101)
    data, true_params = make_psf_model_image(shape, psf_model, n_sources,
                                             model_shape=model_shape,
                                             flux=(500, 700), fwhm=(2.7, 2.7),
                                             min_separation=10, seed=0)
    return data, true_params


@pytest.mark.parametrize('fix_fwhm', [False, True])
def test_fit_2dgaussian_single(fix_fwhm):
    yy, xx = np.mgrid[:51, :51]
    fwhm = 3.123
    model = CircularGaussianPRF(x_0=22.17, y_0=28.87, fwhm=fwhm)
    data = model(xx, yy)

    match = ('fit_shape is None, so the input data array is assumed to be '
             'a cutout image containing only one source')
    with pytest.warns(AstropyUserWarning, match=match):
        fit = fit_2dgaussian(data, fwhm=3, fix_fwhm=fix_fwhm)
    fit_tbl = fit.results
    assert isinstance(fit_tbl, QTable)
    assert len(fit_tbl) == 1
    if fix_fwhm:
        assert 'fwhm_fit' not in fit_tbl.colnames
    else:
        assert 'fwhm_fit' in fit_tbl.colnames
        assert_allclose(fit_tbl['fwhm_fit'], fwhm)

    # test with NaNs - will emit two warnings
    data[22, 29] = np.nan
    with pytest.warns(AstropyUserWarning) as record:
        fit = fit_2dgaussian(data, fwhm=3, fix_fwhm=fix_fwhm)
    # Check that both warnings were emitted
    assert len(record) == 2
    warning_msgs = [str(w.message) for w in record]
    assert any('fit_shape is None' in msg for msg in warning_msgs)
    assert any('non-finite values' in msg for msg in warning_msgs)
    fit_tbl = fit.results
    assert isinstance(fit_tbl, QTable)
    assert len(fit_tbl) == 1

    # test with NaNs and mask - only fit_shape warning
    data[22, 29] = np.nan
    mask = np.isnan(data)
    match = ('fit_shape is None, so the input data array is assumed to be '
             'a cutout image containing only one source')
    with pytest.warns(AstropyUserWarning, match=match):
        fit = fit_2dgaussian(data, fwhm=3, fix_fwhm=fix_fwhm, mask=mask)
    fit_tbl = fit.results
    assert isinstance(fit_tbl, QTable)
    assert len(fit_tbl) == 1


@pytest.mark.parametrize(('fix_fwhm', 'with_units'),
                         [(False, True), (True, False)])
def test_fit_2dgaussian_multiple(test_data, fix_fwhm, with_units):
    data, sources = test_data

    unit = u.nJy
    if with_units:
        data = data * unit

    xypos = list(zip(sources['x_0'], sources['y_0'], strict=True))
    fit = fit_2dgaussian(data, xypos=xypos, fit_shape=(5, 5),
                         fix_fwhm=fix_fwhm)
    fit_tbl = fit.results
    assert isinstance(fit_tbl, QTable)
    assert len(fit_tbl) == len(sources)
    if fix_fwhm:
        assert 'fwhm_fit' not in fit_tbl.colnames
    else:
        assert 'fwhm_fit' in fit_tbl.colnames
        assert_allclose(fit_tbl['fwhm_fit'], sources['fwhm'])

    # test with zip instead of list
    xypos = zip(sources['x_0'], sources['y_0'], strict=True)
    fit2 = fit_2dgaussian(data, xypos=xypos, fit_shape=(5, 5),
                          fix_fwhm=fix_fwhm)
    fit_tbl2 = fit2.results
    assert isinstance(fit_tbl2, QTable)
    assert fit_tbl2.values_equal(fit_tbl)

    if with_units:
        for column in fit_tbl.colnames:
            if 'flux' in column:
                assert fit_tbl['flux_fit'].unit == unit


def test_fit_2dgaussian_fit_shape_none_single():
    """
    Test fit_2dgaussian with fit_shape=None for a single source.

    Should emit a warning and use the data shape (adjusted to odd).
    """
    yy, xx = np.mgrid[:50, :50]  # Even shape
    fwhm = 3.0
    model = CircularGaussianPRF(x_0=25.0, y_0=25.0, fwhm=fwhm)
    data = model(xx, yy)

    match = ('fit_shape is None, so the input data array is assumed to be '
             'a cutout image containing only one source')
    with pytest.warns(AstropyUserWarning, match=match):
        fit = fit_2dgaussian(data, fit_shape=None, fix_fwhm=False)

    fit_tbl = fit.results
    assert isinstance(fit_tbl, QTable)
    assert len(fit_tbl) == 1
    assert_allclose(fit_tbl['fwhm_fit'], fwhm, rtol=0.01)


def test_fit_2dgaussian_fit_shape_none_single_odd():
    """
    Test fit_2dgaussian with fit_shape=None for single source with odd
    shape.

    Should emit a warning and use the data shape unchanged.
    """
    yy, xx = np.mgrid[:51, :51]  # Odd shape
    fwhm = 3.123
    model = CircularGaussianPRF(x_0=25.0, y_0=25.0, fwhm=fwhm)
    data = model(xx, yy)

    match = ('fit_shape is None, so the input data array is assumed to be '
             'a cutout image containing only one source')
    with pytest.warns(AstropyUserWarning, match=match):
        fit = fit_2dgaussian(data, fit_shape=None, fix_fwhm=False)

    fit_tbl = fit.results
    assert isinstance(fit_tbl, QTable)
    assert len(fit_tbl) == 1
    assert_allclose(fit_tbl['fwhm_fit'], fwhm, rtol=0.01)


def test_fit_2dgaussian_fit_shape_none_multiple():
    """
    Test fit_2dgaussian raises ValueError with fit_shape=None for
    multiple sources.
    """
    yy, xx = np.mgrid[:51, :51]
    model1 = CircularGaussianPRF(x_0=20.0, y_0=20.0, fwhm=3.0, flux=100)
    model2 = CircularGaussianPRF(x_0=30.0, y_0=30.0, fwhm=3.0, flux=100)
    data = model1(xx, yy) + model2(xx, yy)

    xypos = [(20.0, 20.0), (30.0, 30.0)]

    match = ('fit_shape is required when fitting multiple sources')
    with pytest.raises(ValueError, match=match):
        fit_2dgaussian(data, xypos=xypos, fit_shape=None)


def test_fit_fwhm_single():
    yy, xx = np.mgrid[:51, :51]
    fwhm0 = 3.123
    model = CircularGaussianPRF(x_0=22.17, y_0=28.87, fwhm=fwhm0)
    data = model(xx, yy)

    # fit_fwhm re-emits the fit_shape warning
    match = ('fit_shape is None, so the input data array is assumed to be '
             'a cutout image containing only one source')
    with pytest.warns(AstropyUserWarning, match=match):
        fwhm = fit_fwhm(data, fwhm=3)
    assert isinstance(fwhm, np.ndarray)
    assert len(fwhm) == 1
    assert_allclose(fwhm, fwhm0)

    # test warning message for convergence issues - flat data should also
    # emit the fit_shape warning since fit_shape=None
    with pytest.warns(AstropyUserWarning) as record:
        fwhm = fit_fwhm(np.zeros(data.shape) + 1)
    # Should get at least the fit_shape warning
    assert len(record) >= 1
    warning_msgs = [str(w.message) for w in record]
    assert any('fit_shape is None' in msg for msg in warning_msgs)
    assert len(fwhm) == 1


@pytest.mark.parametrize('with_units', [False, True])
def test_fit_fwhm_multiple(test_data, with_units):
    data, sources = test_data

    unit = u.nJy
    if with_units:
        data = data * unit

    xypos = list(zip(sources['x_0'], sources['y_0'], strict=True))
    fwhms = fit_fwhm(data, xypos=xypos, fit_shape=(5, 5))
    assert isinstance(fwhms, np.ndarray)
    assert len(fwhms) == len(sources)
    assert_allclose(fwhms, sources['fwhm'])


def test_fit_fwhm_fit_shape_none(test_data):
    data, sources = test_data
    xypos = (sources['x_0'][0], sources['y_0'][0])

    match = ('fit_shape is None, so the input data array is assumed to be '
             'a cutout image')
    with pytest.warns(AstropyUserWarning, match=match):
        fit_fwhm(data, xypos=xypos, fit_shape=None)


def test_interpolate_missing_data():
    data = np.arange(100).reshape(10, 10)
    mask = np.zeros_like(data, dtype=bool)
    mask[5, 5] = True

    data_int = _interpolate_missing_data(data, mask, method='nearest')
    assert 54 <= data_int[5, 5] <= 56

    data_int = _interpolate_missing_data(data, mask, method='cubic')
    assert 54 <= data_int[5, 5] <= 56

    match = 'data must be a 2D array'
    with pytest.raises(ValueError, match=match):
        _interpolate_missing_data(np.arange(10), mask)

    match = 'mask and data must have the same shape'
    with pytest.raises(ValueError, match=match):
        _interpolate_missing_data(data, mask[1:, :])

    match = 'Unsupported interpolation method'
    with pytest.raises(ValueError, match=match):
        _interpolate_missing_data(data, mask, method='invalid')


def test_interpolate_missing_data_edge_pixels():
    """
    Test that edge pixels are always filled with cubic interpolation.
    """
    data = np.arange(100, dtype=float).reshape(10, 10)
    mask = np.zeros_like(data, dtype=bool)

    # Mask corner and edge pixels where cubic interpolation typically
    # fails
    mask[0, 0] = True  # corner
    mask[0, 5] = True  # top edge
    mask[9, 9] = True  # corner
    mask[5, 9] = True  # right edge

    data_int = _interpolate_missing_data(data, mask, method='cubic')

    # All masked pixels should be filled (no NaN values)
    assert np.all(np.isfinite(data_int))
    assert not np.any(np.isnan(data_int[mask]))


def test_interpolate_missing_data_no_mask():
    """
    Test that data is returned unchanged when no pixels are masked.
    """
    data = np.arange(100, dtype=float).reshape(10, 10)
    mask = np.zeros_like(data, dtype=bool)

    data_int = _interpolate_missing_data(data, mask, method='cubic')
    assert np.array_equal(data, data_int)


def test_interpolate_missing_data_all_masked():
    """
    Test that all-masked data returns NaN array.
    """
    data = np.arange(100, dtype=float).reshape(10, 10)
    mask = np.ones_like(data, dtype=bool)  # All pixels masked

    data_int = _interpolate_missing_data(data, mask, method='cubic')

    # All values should be NaN when all data is masked
    assert np.all(np.isnan(data_int))

    # Same for nearest-neighbor method
    data_int = _interpolate_missing_data(data, mask, method='nearest')
    assert np.all(np.isnan(data_int))


def test_validate_psf_model():
    model = np.arange(10)

    match = 'psf_model must be an Astropy Model subclass'
    with pytest.raises(TypeError, match=match):
        _validate_psf_model(model)

    match = 'psf_model must be two-dimensional'
    model = Gaussian1D()
    with pytest.raises(ValueError, match=match):
        _validate_psf_model(model)

    match = 'psf_model must be two-dimensional'
    model = Gaussian1D()
    with pytest.raises(ValueError, match=match):
        _validate_psf_model(model)


def test_get_psf_model_main_params():
    model = CircularGaussianPRF(fwhm=1.0)
    params = _get_psf_model_main_params(model)
    assert len(params) == 3
    assert params == ('x_0', 'y_0', 'flux')

    match = 'Invalid PSF model - could not find PSF parameter names'
    model = Gaussian2D()
    with pytest.raises(ValueError, match=match):
        _get_psf_model_main_params(model)

    set_params = ('x_mean', 'y_mean', 'amplitude')
    model.x_name = set_params[0]
    model.y_name = set_params[1]
    model.flux_name = set_params[2]
    params = _get_psf_model_main_params(model)
    assert len(params) == 3
    assert params == set_params
