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

import itertools

import numpy as np
import pytest
from astropy.modeling.fitting import TRFLSQFitter
from astropy.nddata import (InverseVariance, NDData, StdDevUncertainty,
                            VarianceUncertainty)
from astropy.stats import SigmaClip
from astropy.table import Table
from astropy.utils.exceptions import AstropyUserWarning
from numpy.testing import assert_allclose

from photutils.datasets import make_model_image
from photutils.psf import CircularGaussianPRF, make_psf_model_image
from photutils.psf.epsf import EPSFBuilder, EPSFFitter
from photutils.psf.epsf_stars import EPSFStars, extract_stars


@pytest.fixture
def epsf_test_data():
    """
    Create a simulated image for testing.
    """
    fwhm = 2.7
    psf_model = CircularGaussianPRF(flux=1, fwhm=fwhm)
    model_shape = (9, 9)
    n_sources = 100
    shape = (750, 750)
    data, true_params = make_psf_model_image(shape, psf_model, n_sources,
                                             model_shape=model_shape,
                                             flux=(500, 700),
                                             min_separation=25,
                                             border_size=25, seed=0)

    nddata = NDData(data)
    init_stars = Table()
    init_stars['x'] = true_params['x_0'].astype(int)
    init_stars['y'] = true_params['y_0'].astype(int)

    return {
        'fwhm': fwhm,
        'data': data,
        'nddata': nddata,
        'init_stars': init_stars,
    }


class TestEPSFBuild:

    def test_extract_stars(self, epsf_test_data):
        size = 25
        stars = extract_stars(epsf_test_data['nddata'],
                              epsf_test_data['init_stars'],
                              size=size)

        assert len(stars) == len(epsf_test_data['init_stars'])
        assert isinstance(stars, EPSFStars)
        assert isinstance(stars[0], EPSFStars)
        assert stars[0].data.shape == (size, size)

    def test_extract_stars_uncertainties(self, epsf_test_data):
        rng = np.random.default_rng(0)
        shape = epsf_test_data['nddata'].data.shape
        error = np.abs(rng.normal(loc=0, scale=1, size=shape))
        uncertainty1 = StdDevUncertainty(error)
        uncertainty2 = uncertainty1.represent_as(VarianceUncertainty)
        uncertainty3 = uncertainty1.represent_as(InverseVariance)
        ndd1 = NDData(epsf_test_data['nddata'].data, uncertainty=uncertainty1)
        ndd2 = NDData(epsf_test_data['nddata'].data, uncertainty=uncertainty2)
        ndd3 = NDData(epsf_test_data['nddata'].data, uncertainty=uncertainty3)

        size = 25
        match = 'were not extracted because their cutout region extended'
        ndd_inputs = (ndd1, ndd2, ndd3)

        outputs = [extract_stars(ndd_input, epsf_test_data['init_stars'],
                                 size=size) for ndd_input in ndd_inputs]

        for stars in outputs:
            assert len(stars) == len(epsf_test_data['init_stars'])
            assert isinstance(stars, EPSFStars)
            assert isinstance(stars[0], EPSFStars)
            assert stars[0].data.shape == (size, size)
            assert stars[0].weights.shape == (size, size)

        assert_allclose(outputs[0].weights, outputs[1].weights)
        assert_allclose(outputs[0].weights, outputs[2].weights)

        uncertainty = StdDevUncertainty(np.zeros(shape))
        ndd = NDData(epsf_test_data['nddata'].data, uncertainty=uncertainty)

        match = 'One or more weight values is not finite'
        with pytest.warns(AstropyUserWarning, match=match):
            stars = extract_stars(ndd, epsf_test_data['init_stars'][0:3],
                                  size=size)

    @pytest.mark.parametrize('shape', [(25, 25), (19, 25), (25, 19)])
    def test_epsf_build(self, epsf_test_data, shape):
        """
        This is an end-to-end test of EPSFBuilder on a simulated image.
        """
        oversampling = 2
        stars = extract_stars(epsf_test_data['nddata'],
                              epsf_test_data['init_stars'][:10],
                              size=shape)
        epsf_builder = EPSFBuilder(oversampling=oversampling, maxiters=5,
                                   progress_bar=False, norm_radius=10,
                                   recentering_maxiters=5)
        epsf, fitted_stars = epsf_builder(stars)

        ref_size = np.array(shape) * oversampling + 1
        assert epsf.data.shape == tuple(ref_size)

        # Verify basic EPSF properties
        assert len(fitted_stars) == 10
        assert epsf.data.sum() > 2  # Check it has reasonable total flux
        assert epsf.data.max() > 0.01  # Should have a peak

        # Check that the center region has higher values than edges
        center_y, center_x = np.array(ref_size) // 2
        center_val = epsf.data[center_y, center_x]
        edge_val = epsf.data[0, 0]
        assert center_val > edge_val  # Center should be brighter than edge

        # Test that residual computation works (basic functionality test)
        resid_star = fitted_stars[0].compute_residual_image(epsf)
        assert isinstance(resid_star, np.ndarray)
        assert resid_star.shape == fitted_stars[0].data.shape

    def test_epsf_fitting_bounds(self, epsf_test_data):
        size = 25
        oversampling = 4
        stars = extract_stars(epsf_test_data['nddata'],
                              epsf_test_data['init_stars'],
                              size=size)

        epsf_builder = EPSFBuilder(oversampling=oversampling, maxiters=8,
                                   progress_bar=True, norm_radius=25,
                                   recentering_maxiters=5,
                                   fitter=EPSFFitter(fit_boxsize=31),
                                   smoothing_kernel='quadratic')

        # With a boxsize larger than the cutout we expect the fitting to
        # fail for all stars, due to star._fit_error_status
        match1 = 'The ePSF fitting failed for all stars'
        match2 = r'The star at .* cannot be fit because its fitting region '
        with (pytest.raises(ValueError, match=match1),
                pytest.warns(AstropyUserWarning, match=match2)):
            epsf_builder(stars)

    def test_epsf_build_invalid_fitter(self):
        """
        Test that the input fitter is an EPSFFitter instance.
        """
        match = 'fitter must be an EPSFFitter instance'
        with pytest.raises(TypeError, match=match):
            EPSFBuilder(fitter=EPSFFitter, maxiters=3)

        with pytest.raises(TypeError, match=match):
            EPSFBuilder(fitter=TRFLSQFitter(), maxiters=3)

        with pytest.raises(TypeError, match=match):
            EPSFBuilder(fitter=TRFLSQFitter, maxiters=3)


def test_epsfbuilder_inputs():
    # invalid inputs
    match = "'oversampling' must be specified"
    with pytest.raises(ValueError, match=match):
        EPSFBuilder(oversampling=None)
    match = 'oversampling must be > 0'
    with pytest.raises(ValueError, match=match):
        EPSFBuilder(oversampling=-1)
    match = 'maxiters must be a positive number'
    with pytest.raises(ValueError, match=match):
        EPSFBuilder(maxiters=-1)
    match = 'oversampling must be > 0'
    with pytest.raises(ValueError, match=match):
        EPSFBuilder(oversampling=[-1, 4])

    # valid inputs
    EPSFBuilder(oversampling=6)
    EPSFBuilder(oversampling=[4, 6])

    # invalid inputs
    for sigma_clip in [None, [], 'a']:
        match = 'sigma_clip must be an astropy.stats.SigmaClip instance'
        with pytest.raises(TypeError, match=match):
            EPSFBuilder(sigma_clip=sigma_clip)

    # valid inputs
    EPSFBuilder(sigma_clip=SigmaClip(sigma=2.5, cenfunc='mean', maxiters=2))


@pytest.mark.parametrize('oversamp', [3, 4])
def test_epsf_build_oversampling(oversamp):
    offsets = (np.arange(oversamp) * 1.0 / oversamp - 0.5 + 1.0
               / (2.0 * oversamp))
    xydithers = np.array(list(itertools.product(offsets, offsets)))
    xdithers = np.transpose(xydithers)[0]
    ydithers = np.transpose(xydithers)[1]

    nstars = oversamp**2
    fwhm = 7.0
    sources = Table()
    offset = 50
    size = oversamp * offset + offset
    y, x = np.mgrid[0:oversamp, 0:oversamp] * offset + offset
    sources['x_0'] = x.ravel() + xdithers
    sources['y_0'] = y.ravel() + ydithers
    sources['fwhm'] = np.full((nstars,), fwhm)

    psf_model = CircularGaussianPRF(fwhm=fwhm)
    shape = (size, size)
    data = make_model_image(shape, psf_model, sources)
    nddata = NDData(data=data)
    stars_tbl = Table()
    stars_tbl['x'] = sources['x_0']
    stars_tbl['y'] = sources['y_0']
    stars = extract_stars(nddata, stars_tbl, size=25)
    epsf_builder = EPSFBuilder(oversampling=oversamp, maxiters=15,
                               progress_bar=False, recentering_maxiters=20)
    epsf, _ = epsf_builder(stars)

    # input PSF shape
    size = epsf.data.shape[0]
    cen = (size - 1) / 2
    fwhm2 = oversamp * fwhm
    m = CircularGaussianPRF(flux=1, x_0=cen, y_0=cen, fwhm=fwhm2)
    yy, xx = np.mgrid[0:size, 0:size]
    psf = m(xx, yy)

    assert_allclose(epsf.data, psf * epsf.data.sum(), atol=2.5e-4)
