# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""Tests for polynomial models."""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import os

from itertools import product

import numpy as np

from numpy.testing.utils import assert_allclose

from .. import fitting
from ...tests.helper import pytest
from ... import wcs
from ...io import fits
from ..polynomial import (Chebyshev1D, Legendre1D, Polynomial1D,
                          Chebyshev2D, Legendre2D, Polynomial2D, SIP,
                          PolynomialBase, OrthoPolynomialBase)
from ..functional_models import Linear1D
from ...utils.data import get_pkg_data_filename

try:
    from scipy import optimize  # pylint: disable=W0611
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False


linear1d = {
    Chebyshev1D: {
        'args': (3,),
        'kwargs': {'domain': [1, 10]},
        'parameters': {'c0': 1.2, 'c1': 2, 'c2': 2.3, 'c3': 0.2},
        'constraints': {'fixed': {'c0': 1.2}}
    },
    Legendre1D: {
        'args': (3,),
        'kwargs': {'domain': [1, 10]},
        'parameters': {'c0': 1.2, 'c1': 2, 'c2': 2.3, 'c3': 0.2},
        'constraints': {'fixed': {'c0': 1.2}}
    },
    Polynomial1D: {
        'args': (3,),
        'kwargs': {'domain': [1, 10]},
        'parameters': {'c0': 1.2, 'c1': 2, 'c2': 2.3, 'c3': 0.2},
        'constraints': {'fixed': {'c0': 1.2}}
    },
    Linear1D: {
        'args': (),
        'kwargs': {},
        'parameters': {'intercept': 1.2, 'slope': 23.1},
        'constraints': {'fixed': {'intercept': 1.2}}
    }
}


linear2d = {
    Chebyshev2D: {
        'args': (1, 1),
        'kwargs': {'x_domain': [0, 99], 'y_domain': [0, 82]},
        'parameters': {'c0_0': 1.2, 'c1_0': 2, 'c0_1': 2.3, 'c1_1': 0.2},
        'constraints': {'fixed': {'c0_0': 1.2}}
    },
    Legendre2D: {
        'args': (1, 1),
        'kwargs': {'x_domain': [0, 99], 'y_domain': [0, 82]},
        'parameters': {'c0_0': 1.2, 'c1_0': 2, 'c0_1': 2.3, 'c1_1': 0.2},
        'constraints': {'fixed': {'c0_0': 1.2}}
    },
    Polynomial2D: {
        'args': (1,),
        'kwargs': {},
        'parameters': {'c0_0': 1.2, 'c1_0': 2, 'c0_1': 2.3},
        'constraints': {'fixed': {'c0_0': 1.2}}
    }
}


@pytest.mark.skipif('not HAS_SCIPY')
class TestFitting(object):
    """Test linear fitter with polynomial models."""

    def setup_class(self):
        self.N = 100
        self.M = 100
        self.x1 = np.linspace(1, 10, 100)
        self.y2, self.x2 = np.mgrid[:100, :83]
        rsn = np.random.RandomState(0)
        self.n1 = rsn.randn(self.x1.size) * .1
        self.n2 = rsn.randn(self.x2.size)
        self.n2.shape = self.x2.shape
        self.linear_fitter = fitting.LinearLSQFitter()
        self.non_linear_fitter = fitting.LevMarLSQFitter()

    # TODO: Most of these test cases have some pretty repetitive setup that we
    # could probably factor out

    @pytest.mark.parametrize(('model_class','constraints'),
                             list(product(linear1d.keys(), (False, True))))
    def test_linear_fitter_1D(self, model_class, constraints):
        """Test fitting with LinearLSQFitter"""

        model_args = linear1d[model_class]
        kwargs = {}
        kwargs.update(model_args['kwargs'])
        kwargs.update(model_args['parameters'])

        if constraints:
            kwargs.update(model_args['constraints'])

        model = model_class(*model_args['args'], **kwargs)

        y1 = model(self.x1)
        model_lin = self.linear_fitter(model, self.x1, y1 + self.n1)

        if constraints:
            # For the constraints tests we're not checking the overall fit,
            # just that the constraint was maintained
            fixed = model_args['constraints'].get('fixed', None)
            if fixed:
                for param, value in fixed.items():
                    expected = model_args['parameters'][param]
                    assert getattr(model_lin, param).value == expected
        else:
            assert_allclose(model_lin.parameters, model.parameters,
                            atol=0.2)

    @pytest.mark.parametrize(('model_class','constraints'),
                             list(product(linear1d.keys(), (False, True))))
    def test_non_linear_fitter_1D(self, model_class, constraints):
        """Test fitting with non-linear LevMarLSQFitter"""

        model_args = linear1d[model_class]
        kwargs = {}
        kwargs.update(model_args['kwargs'])
        kwargs.update(model_args['parameters'])

        if constraints:
            kwargs.update(model_args['constraints'])

        model = model_class(*model_args['args'], **kwargs)

        y1 = model(self.x1)
        model_nlin = self.non_linear_fitter(model, self.x1, y1 + self.n1)

        if constraints:
            fixed = model_args['constraints'].get('fixed', None)
            if fixed:
                for param, value in fixed.items():
                    expected = model_args['parameters'][param]
                    assert getattr(model_nlin, param).value == expected
        else:
            assert_allclose(model_nlin.parameters, model.parameters,
                            atol=0.2)

    @pytest.mark.parametrize(('model_class','constraints'),
                             list(product(linear2d.keys(), (False, True))))
    def test_linear_fitter_2D(self, model_class, constraints):
        """Test fitting with LinearLSQFitter"""

        model_args = linear2d[model_class]
        kwargs = {}
        kwargs.update(model_args['kwargs'])
        kwargs.update(model_args['parameters'])

        if constraints:
            kwargs.update(model_args['constraints'])

        model = model_class(*model_args['args'], **kwargs)

        z = model(self.x2, self.y2)
        model_lin = self.linear_fitter(model, self.x2, self.y2, z + self.n2)

        if constraints:
            fixed = model_args['constraints'].get('fixed', None)
            if fixed:
                for param, value in fixed.items():
                    expected = model_args['parameters'][param]
                    assert getattr(model_lin, param).value == expected
        else:
            assert_allclose(model_lin.parameters, model.parameters,
                            atol=0.2)

    @pytest.mark.parametrize(('model_class','constraints'),
                             list(product(linear2d.keys(), (False, True))))
    def test_non_linear_fitter_2D(self, model_class, constraints):
        """Test fitting with non-linear LevMarLSQFitter"""

        model_args = linear2d[model_class]
        kwargs = {}
        kwargs.update(model_args['kwargs'])
        kwargs.update(model_args['parameters'])

        if constraints:
            kwargs.update(model_args['constraints'])

        model = model_class(*model_args['args'], **kwargs)

        z = model(self.x2, self.y2)
        model_nlin = self.non_linear_fitter(model, self.x2, self.y2,
                                            z + self.n2)

        if constraints:
            fixed = model_args['constraints'].get('fixed', None)
            if fixed:
                for param, value in fixed.items():
                    expected = model_args['parameters'][param]
                    assert getattr(model_nlin, param).value == expected
        else:
            assert_allclose(model_nlin.parameters, model.parameters,
                            atol=0.2)


@pytest.mark.parametrize('model_class',
                         [cls for cls in list(linear1d) + list(linear2d)
                          if isinstance(cls, PolynomialBase)])
def test_polynomial_init_with_constraints(model_class):
    """
    Test that polynomial models can be instantiated with constraints, but no
    parameters specified.

    Regression test for https://github.com/astropy/astropy/issues/3606
    """

    # Just determine which parameter to place a constraint on; it doesn't
    # matter which parameter it is to exhibit the problem so long as it's a
    # valid parameter for the model
    if '1D' in model_class.__name__:
        param = 'c0'
    else:
        param = 'c0_0'

    if issubclass(model_class, OrthoPolynomialBase):
        degree = (2, 2)
    else:
        degree = (2,)

    m = model_class(*degree, fixed={param: True})

    assert m.fixed[param] is True
    assert getattr(m, param).fixed is True


def test_sip_hst():
    """Test SIP against astropy.wcs"""

    test_file = get_pkg_data_filename(os.path.join('data', 'hst_sip.hdr'))
    hdr = fits.Header.fromtextfile(test_file)
    crpix1 = hdr['CRPIX1']
    crpix2 = hdr['CRPIX2']
    wobj = wcs.WCS(hdr)
    a_pars = dict(**hdr['A_*'])
    b_pars = dict(**hdr['B_*'])
    a_order = a_pars.pop('A_ORDER')
    b_order = b_pars.pop('B_ORDER')
    sip = SIP([crpix1, crpix2], a_order, b_order, a_pars, b_pars)
    coords = [1, 1]
    rel_coords = [1 - crpix1, 1 - crpix2]
    astwcs_result = wobj.sip_pix2foc([coords], 1)[0] - rel_coords
    assert_allclose(sip(1, 1), astwcs_result)


def test_sip_irac():
    """Test forward and inverse SIP againts astropy.wcs"""

    test_file = get_pkg_data_filename(os.path.join('data', 'irac_sip.hdr'))
    hdr = fits.Header.fromtextfile(test_file)
    crpix1 = hdr['CRPIX1']
    crpix2 = hdr['CRPIX2']
    wobj = wcs.WCS(hdr)
    a_pars = dict(**hdr['A_*'])
    b_pars = dict(**hdr['B_*'])
    ap_pars = dict(**hdr['AP_*'])
    bp_pars = dict(**hdr['BP_*'])
    a_order = a_pars.pop('A_ORDER')
    b_order = b_pars.pop('B_ORDER')
    ap_order = ap_pars.pop('AP_ORDER')
    bp_order = bp_pars.pop('BP_ORDER')
    del a_pars['A_DMAX']
    del b_pars['B_DMAX']
    pix = [200, 200]
    rel_pix = [200 - crpix1, 200 - crpix2]
    sip = SIP([crpix1, crpix2], a_order, b_order, a_pars, b_pars,
              ap_order=ap_order, ap_coeff=ap_pars, bp_order=bp_order,
              bp_coeff=bp_pars)

    foc = wobj.sip_pix2foc([pix], 1)
    newpix = wobj.sip_foc2pix(foc, 1)[0]
    assert_allclose(sip(*pix), foc[0] - rel_pix)
    assert_allclose(sip.inverse(*foc[0]) +
                    foc[0] - rel_pix, newpix - pix)


def test_sip_no_coeff():
    sip = SIP([10,12], 2, 2)
    assert_allclose(sip.sip1d_a.parameters, [0., 0., 0])
    assert_allclose(sip.sip1d_b.parameters, [0., 0., 0])
    with pytest.raises(NotImplementedError):
        sip.inverse


@pytest.mark.parametrize('cls', (Polynomial1D, Chebyshev1D, Legendre1D,
                                 Polynomial2D, Chebyshev2D, Legendre2D))
def test_zero_degree_polynomial(cls):
    """
    A few tests that degree=0 polynomials are correctly evaluated and
    fitted.

    Regression test for https://github.com/astropy/astropy/pull/3589
    """

    if cls.n_inputs == 1:  # Test 1D polynomials
        p1 = cls(degree=0, c0=1)
        assert p1(0) == 1
        assert np.all(p1(np.zeros(5)) == np.ones(5))

        x = np.linspace(0, 1, 100)
        # Add a little noise along a straight line
        y = 1 + np.random.uniform(0, 0.1, len(x))

        p1_init = cls(degree=0)
        fitter = fitting.LinearLSQFitter()
        p1_fit = fitter(p1_init, x, y)

        # The fit won't be exact of course, but it should get close to within
        # 1%
        assert_allclose(p1_fit.c0, 1, atol=0.10)
    elif cls.n_inputs == 2:  # Test 2D polynomials
        if issubclass(cls, OrthoPolynomialBase):
            p2 = cls(x_degree=0, y_degree=0, c0_0=1)
        else:
            p2 = cls(degree=0, c0_0=1)
        assert p2(0, 0) == 1
        assert np.all(p2(np.zeros(5), np.zeros(5)) == np.ones(5))

        y, x = np.mgrid[0:1:100j,0:1:100j]
        z = (1 + np.random.uniform(0, 0.1, x.size)).reshape(100, 100)

        if issubclass(cls, OrthoPolynomialBase):
            p2_init = cls(x_degree=0, y_degree=0)
        else:
            p2_init = cls(degree=0)
        fitter = fitting.LinearLSQFitter()
        p2_fit = fitter(p2_init, x, y, z)

        assert_allclose(p2_fit.c0_0, 1, atol=0.10)
