"""Test the features of the lightkurve.prf.prfmodels module."""
from __future__ import division, print_function
from collections import OrderedDict

from astropy.io import fits
from astropy.utils.data import get_pkg_data_filename
import numpy as np
from numpy.testing import assert_allclose
import pytest

from lightkurve.prf import KeplerPRF, SimpleKeplerPRF
from lightkurve.targetpixelfile import KeplerTargetPixelFile


@pytest.mark.remote_data  # PRF relies on calibration files on stsci.edu
def test_prf_normalization():
    """Does the PRF model integrate to the requested flux across the focal plane?"""
    for channel in [1, 20, 40, 60, 84]:
        for col in [123, 678]:
            for row in [234, 789]:
                shape = (18, 14)
                flux = 100
                prf = KeplerPRF(channel=channel, column=col, row=row, shape=shape)
                prf_sum = prf.evaluate(
                    col + shape[0] / 2, row + shape[1] / 2, flux, 1, 1, 0
                ).sum()
                assert np.isclose(prf_sum, flux, rtol=0.1)


@pytest.mark.remote_data  # PRF relies on calibration files on stsci.edu
def test_simple_kepler_prf():
    """Ensures that concentric PRFs have the same values."""
    prf_1 = SimpleKeplerPRF(channel=16, shape=[20, 20], column=0, row=0)
    prf_2 = SimpleKeplerPRF(channel=16, shape=[10, 10], column=5, row=5)
    for c in [10, 8, 10, 7]:
        for r in [10, 10, 7, 7]:
            assert_allclose(
                prf_2(center_col=c, center_row=r, flux=1),
                prf_1(center_col=c, center_row=r, flux=1)[5:15, 5:15],
                rtol=1e-5,
            )


@pytest.mark.remote_data
def test_simple_kepler_prf_interpolation_consistency():
    """Ensures that the interpolated prf is consistent with calibration files."""
    sprf = SimpleKeplerPRF(channel=56, shape=[15, 15], column=0, row=0)
    cal_prf = fits.open(
        "http://archive.stsci.edu/missions/kepler/fpc/prf/" "kplr16.4_2011265_prf.fits"
    )
    cal_prf_subsampled = cal_prf[-1].data[25::50, 25::50]
    cal_prf_subsampled_normalized = cal_prf_subsampled / (
        cal_prf[-1].data.sum() * 0.02 ** 2
    )
    sprf_data = sprf(center_col=7.5, center_row=7.5, flux=1)
    np.isclose(np.sum(np.abs(sprf_data - cal_prf_subsampled_normalized)), 0)


@pytest.mark.remote_data  # PRF relies on calibration files on stsci.edu
def test_get_model_prf():
    tpf_fn = get_pkg_data_filename("../../tests/data/test-tpf-star.fits")
    tpf = KeplerTargetPixelFile(tpf_fn)

    prf = KeplerPRF(
        channel=tpf.channel, shape=tpf.shape[1:], column=tpf.column, row=tpf.row
    )
    prf_from_tpf = tpf.get_prf_model()

    assert type(prf) == type(prf_from_tpf)
    assert prf.channel == prf_from_tpf.channel
    assert prf.shape == prf_from_tpf.shape
    assert prf.column == prf_from_tpf.column
    assert prf.row == prf_from_tpf.row


@pytest.mark.remote_data  # PRF relies on calibration files on stsci.edu
def test_keplerprf_gradient_against_simplekeplerprf():
    """is the gradient of KeplerPRF consistent with
    the gradient of SimpleKeplerPRF?
    """
    kwargs = {"channel": 56, "shape": [15, 15], "column": 0, "row": 0}
    params = {"center_col": 7, "center_row": 7, "flux": 1.0}
    simple_prf = SimpleKeplerPRF(**kwargs)
    prf = KeplerPRF(**kwargs)
    prf_grad = prf.gradient(rotation_angle=0.0, scale_col=1.0, scale_row=1.0, **params)
    assert_allclose(prf_grad[:-3], simple_prf.gradient(**params))


@pytest.mark.remote_data  # PRF relies on calibration files on stsci.edu
@pytest.mark.parametrize(
    "param_to_test",
    [
        ("center_col"),
        ("center_row"),
        ("flux"),
        ("scale_col"),
        ("scale_row"),
        ("rotation_angle"),
    ],
)
def test_keplerprf_gradient_against_calculus(param_to_test):
    """is the gradient of KeplerPRF consistent with Calculus?"""
    params = OrderedDict(
        [
            ("center_col", 7),
            ("center_row", 7),
            ("flux", 1000.0),
            ("scale_col", 1.0),
            ("scale_row", 1.0),
            ("rotation_angle", 0),
        ]
    )
    param_order = OrderedDict(zip(params.keys(), range(0, 6)))
    kwargs = {"channel": 56, "shape": [15, 15], "column": 0, "row": 0}

    prf = KeplerPRF(**kwargs)
    h = 1e-8
    f = prf.evaluate
    inc_params = params.copy()
    # increment the parameter under test for later finite difference computation
    inc_params[param_to_test] += h
    # compute finite differences
    diff_prf = (f(**inc_params) - f(**params)) / h
    # compute analytical gradient
    prf_grad = prf.gradient(**params)
    # assert that the average absolute/relative error is less than 1e-5
    assert (
        np.max(
            np.abs(prf_grad[param_order[param_to_test]] - diff_prf)
            / (1.0 + np.abs(diff_prf))
        )
        < 1e-5
    )
