# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
This module tests model set evaluation and fitting for some common use cases.
"""
# pylint: disable=invalid-name
import pytest
import numpy as np
from numpy.testing import assert_allclose

from astropy.modeling.models import (Polynomial1D, Polynomial2D, Legendre1D, Legendre2D,
                                     Chebyshev2D, Chebyshev1D, Hermite1D, Hermite2D,
                                     Linear1D, Planar2D)
from astropy.modeling.fitting import LinearLSQFitter, FittingWithOutlierRemoval
from astropy.modeling.core import Model
from astropy.modeling.parameters import Parameter
from astropy.utils import NumpyRNGContext
from astropy.stats import sigma_clip


x = np.arange(4)
xx = np.array([x, x + 10])
xxx = np.arange(24).reshape((3, 4, 2))

_RANDOM_SEED = 0x1337


class TParModel(Model):
    """
    A toy model to test parameters machinery
    """
    # standard_broadasting = False
    n_inputs = 1
    outputs = ('x',)
    coeff = Parameter()
    e = Parameter()

    def __init__(self, coeff, e, **kwargs):
        super().__init__(coeff=coeff, e=e, **kwargs)

    @staticmethod
    def evaluate(x, coeff, e):
        return x*coeff + e


@pytest.mark.parametrize('model_class', [Polynomial1D, Chebyshev1D, Legendre1D, Hermite1D])
def test_model1d_axis_1(model_class):
    """
    Test that a model initialized with model_set_axis=1
    can be evaluated with model_set_axis=False.
    """
    n_models = 2
    model_axis = 1

    c0 = [[2, 3]]
    c1 = [[1, 2]]
    t1 = model_class(1, c0=2, c1=1)
    t2 = model_class(1, c0=3, c1=2)

    p1 = model_class(1, c0=c0, c1=c1, n_models=n_models, model_set_axis=model_axis)

    with pytest.raises(ValueError):
        p1(x)

    y = p1(x, model_set_axis=False)
    assert y.shape[model_axis] == n_models
    assert_allclose(y[:, 0], t1(x))
    assert_allclose(y[:, 1], t2(x))

    y = p1(xx, model_set_axis=False)
    assert y.shape[model_axis] == n_models
    assert_allclose(y[:, 0, :], t1(xx))
    assert_allclose(y[:, 1, :], t2(xx))

    y = p1(xxx, model_set_axis=False)
    assert y.shape[model_axis] == n_models
    assert_allclose(y[:, 0, :, :], t1(xxx))
    assert_allclose(y[:, 1, :, :], t2(xxx))


@pytest.mark.parametrize('model_class', [Polynomial1D, Chebyshev1D, Legendre1D, Hermite1D])
def test_model1d_axis_2(model_class):
    """
    Test that a model initialized with model_set_axis=2
    can be evaluated with model_set_axis=False.
    """
    p1 = model_class(1, c0=[[[1, 2, 3]]], c1=[[[10, 20, 30]]],
                     n_models=3, model_set_axis=2)
    t1 = model_class(1, c0=1, c1=10)
    t2 = model_class(1, c0=2, c1=20)
    t3 = model_class(1, c0=3, c1=30)

    with pytest.raises(ValueError):
        p1(x)

    with pytest.raises(ValueError):
        p1(xx)

    y = p1(x, model_set_axis=False)
    assert y.shape == (1, 4, 3)
    assert_allclose(y[:, :, 0].flatten(), t1(x))
    assert_allclose(y[:, :, 1].flatten(), t2(x))
    assert_allclose(y[:, :, 2].flatten(), t3(x))


@pytest.mark.parametrize('model_class', [Polynomial1D, Chebyshev1D, Legendre1D, Hermite1D])
def test_model1d_axis_0(model_class):
    """
    Test that a model initialized with model_set_axis=0
    can be evaluated with model_set_axis=False.
    """
    p1 = model_class(1, n_models=2, model_set_axis=0)
    p1.c0 = [2, 3]
    p1.c1 = [1, 2]
    t1 = model_class(1, c0=2, c1=1)
    t2 = model_class(1, c0=3, c1=2)

    with pytest.raises(ValueError):
        p1(x)

    y = p1(xx)
    assert len(y) == 2
    assert_allclose(y[0], t1(xx[0]))
    assert_allclose(y[1], t2(xx[1]))

    y = p1(x, model_set_axis=False)
    assert len(y) == 2
    assert_allclose(y[0], t1(x))
    assert_allclose(y[1], t2(x))

    y = p1(xx, model_set_axis=False)
    assert len(y) == 2
    assert_allclose(y[0], t1(xx))
    assert_allclose(y[1], t2(xx))
    y = p1(xxx, model_set_axis=False)
    assert_allclose(y[0], t1(xxx))
    assert_allclose(y[1], t2(xxx))
    assert len(y) == 2


@pytest.mark.parametrize('model_class', [Chebyshev2D, Legendre2D, Hermite2D])
def test_model2d_axis_2(model_class):
    """
    Test that a model initialized with model_set_axis=2
    can be evaluated with model_set_axis=False.
    """
    p2 = model_class(1, 1, c0_0=[[[0, 1, 2]]], c0_1=[[[3, 4, 5]]],
                     c1_0=[[[5, 6, 7]]], c1_1=[[[1,1,1]]], n_models=3, model_set_axis=2)
    t1 = model_class(1, 1, c0_0=0, c0_1=3, c1_0=5, c1_1=1)
    t2 = model_class(1, 1, c0_0=1, c0_1=4, c1_0=6, c1_1=1)
    t3 = model_class(1, 1, c0_0=2, c0_1=5, c1_0=7, c1_1=1)

    assert p2.c0_0.shape == (1, 1, 3)
    y = p2(x, x, model_set_axis=False)
    assert y.shape == (1, 4, 3)
    # These are columns along the 2nd axis.
    assert_allclose(y[:, :, 0].flatten(), t1(x, x))
    assert_allclose(y[:, :, 1].flatten(), t2(x, x))
    assert_allclose(y[:, :, 2].flatten(), t3(x, x))


def test_negative_axis():
    p1 = Polynomial1D(1, c0=[1, 2], c1=[3, 4], n_models=2, model_set_axis=-1)
    t1 = Polynomial1D(1, c0=1, c1=3)
    t2 = Polynomial1D(1, c0=2, c1=4)

    with pytest.raises(ValueError):
        p1(x)

    with pytest.raises(ValueError):
        p1(xx)

    xxt = xx.T
    y = p1(xxt)
    assert_allclose(y[:, 0], t1(xxt[:, 0]))
    assert_allclose(y[:, 1], t2(xxt[:, 1]))


def test_shapes():
    p2 = Polynomial1D(1, n_models=3, model_set_axis=2)
    assert p2.c0.shape == (1, 1, 3)
    assert p2.c1.shape == (1, 1, 3)

    p1 = Polynomial1D(1, n_models=2, model_set_axis=1)
    assert p1.c0.shape == (1, 2)
    assert p1.c1.shape == (1, 2)

    p1 = Polynomial1D(1, c0=[1, 2], c1=[3, 4], n_models=2, model_set_axis=-1)
    assert p1.c0.shape == (2,)
    assert p1.c1.shape == (2,)

    e1 = [1, 2]
    e2 = [3, 4]

    a1 = np.array([[10, 20], [30, 40]])
    a2 = np.array([[50, 60], [70, 80]])

    t = TParModel([a1, a2], [e1, e2], n_models=2, model_set_axis=-1)
    assert t.coeff.shape == (2, 2, 2)
    assert t.e.shape == (2, 2)

    t = TParModel([[a1, a2]], [[e1, e2]], n_models=2, model_set_axis=1)
    assert t.coeff.shape == (1, 2, 2, 2)
    assert t.e.shape == (1, 2, 2)

    t = TParModel([a1, a2], [e1, e2], n_models=2, model_set_axis=0)
    assert t.coeff.shape == (2, 2, 2)
    assert t.e.shape == (2, 2)

    t = TParModel([a1, a2], e=[1, 2], n_models=2, model_set_axis=0)
    assert t.coeff.shape == (2, 2, 2)
    assert t.e.shape == (2,)


def test_eval():
    """ Tests evaluation of Linear1D and Planar2D with different model_set_axis."""
    model = Linear1D(slope=[1, 2], intercept=[3, 4], n_models=2)
    p = Polynomial1D(1, c0=[3, 4], c1=[1, 2], n_models=2)

    assert_allclose(model(xx), p(xx))
    assert_allclose(model(x, model_set_axis=False), p(x, model_set_axis=False))

    with pytest.raises(ValueError):
        model(x)

    model = Linear1D(slope=[[1, 2]], intercept=[[3, 4]], n_models=2, model_set_axis=1)
    p = Polynomial1D(1, c0=[[3, 4]], c1=[[1, 2]], n_models=2, model_set_axis=1)

    assert_allclose(model(xx.T), p(xx.T))
    assert_allclose(model(x, model_set_axis=False), p(x, model_set_axis=False))
    with pytest.raises(ValueError):
        model(xx)

    model = Planar2D(slope_x=[1, 2], slope_y=[1, 2], intercept=[3, 4], n_models=2)
    y = model(xx, xx)

    assert y.shape == (2, 4)

    with pytest.raises(ValueError):
        model(x)


# Test fitting


@pytest.mark.parametrize('model_class', [Polynomial1D, Chebyshev1D, Legendre1D, Hermite1D])
def test_linearlsqfitter(model_class):
    """
    Issue #7159
    """
    p = model_class(1, n_models=2, model_set_axis=1)

    # Generate data for fitting 2 models and re-stack them along the last axis:
    y = np.array([2*x+1, x+4])
    y = np.rollaxis(y, 0, -1).T

    f = LinearLSQFitter()
    # This seems to fit the model_set correctly:
    fit = f(p, x, y)
    model_y = fit(x, model_set_axis=False)

    m1 = model_class(1, c0=fit.c0[0][0], c1=fit.c1[0][0], domain=fit.domain)
    m2 = model_class(1, c0=fit.c0[0][1], c1=fit.c1[0][1], domain=fit.domain)
    assert_allclose(model_y[:, 0], m1(x))
    assert_allclose(model_y[:, 1], m2(x))

    p = model_class(1, n_models=2, model_set_axis=0)
    fit = f(p, x, y.T)


def test_model_set_axis_outputs():
    fitter = LinearLSQFitter()
    model_set = Polynomial2D(1, n_models=2, model_set_axis=2)
    y2, x2 = np.mgrid[: 5, : 5]
    # z = np.moveaxis([x2 + y2, 1 - 0.1 * x2 + 0.2 * y2]), 0, 2)
    z = np.rollaxis(np.array([x2 + y2, 1 - 0.1 * x2 + 0.2 * y2]), 0, 3)
    model = fitter(model_set, x2, y2, z)
    res = model(x2, y2, model_set_axis=False)
    assert z.shape == res.shape

    # Test initializing with integer model_set_axis
    # and evaluating with a different model_set_axis
    model_set = Polynomial1D(1, c0=[1, 2], c1=[2, 3],
                             n_models=2, model_set_axis=0)
    y0 = model_set(xx)
    y1 = model_set(xx.T, model_set_axis=1)
    assert_allclose(y0[0], y1[:, 0])
    assert_allclose(y0[1], y1[:, 1])

    model_set = Polynomial1D(1, c0=[[1, 2]], c1=[[2, 3]],
                             n_models=2, model_set_axis=1)
    y0 = model_set(xx.T)
    y1 = model_set(xx, model_set_axis=0)
    assert_allclose(y0[:, 0], y1[0])
    assert_allclose(y0[:, 1], y1[1])
    with pytest.raises(ValueError):
        model_set(x)


def test_fitting_shapes():
    """ Test fitting model sets of Linear1D and Planar2D."""
    fitter = LinearLSQFitter()

    model = Linear1D(slope=[1, 2], intercept=[3, 4], n_models=2)
    y = model(xx)
    fit_model = fitter(model, x, y)

    model = Linear1D(slope=[[1, 2]], intercept=[[3, 4]], n_models=2, model_set_axis=1)
    fit_model = fitter(model, x, y.T)

    model = Planar2D(slope_x=[1, 2], slope_y=[1, 2], intercept=[3, 4], n_models=2)
    y = model(xx, xx)
    fit_model = fitter(model, x, x, y)


def test_compound_model_sets():
    with pytest.raises(ValueError):
        Polynomial1D(1, n_models=2, model_set_axis=1) | Polynomial1D(1, n_models=2, model_set_axis=0)


def test_linear_fit_model_set_errors():
    init_model = Polynomial1D(degree=2, c0=[1, 1], n_models=2)
    x = np.arange(10)
    y = init_model(x, model_set_axis=False)

    fitter = LinearLSQFitter()
    with pytest.raises(ValueError):
        fitter(init_model, x[:5], y)
    with pytest.raises(ValueError):
        fitter(init_model, x, y[:, :5])


def test_linear_fit_model_set_common_weight():
    """Tests fitting multiple models simultaneously."""

    init_model = Polynomial1D(degree=2, c0=[1, 1], n_models=2)
    x = np.arange(10)
    y_expected = init_model(x, model_set_axis=False)
    assert y_expected.shape == (2, 10)

    # Add a bit of random noise
    with NumpyRNGContext(_RANDOM_SEED):
        y = y_expected + np.random.normal(0, 0.01, size=y_expected.shape)

    fitter = LinearLSQFitter()
    weights = np.ones(10)
    weights[[0, -1]] = 0
    fitted_model = fitter(init_model, x, y, weights=weights)
    assert_allclose(fitted_model(x, model_set_axis=False), y_expected,
                    rtol=1e-1)

    # Check that using null weights raises an error
    # ValueError: On entry to DLASCL parameter number 4 had an illegal value
    with pytest.raises(ValueError,
                       match='Found NaNs in the coefficient matrix'):
        with pytest.warns(RuntimeWarning,
                          match=r'invalid value encountered in.*divide'):
            fitted_model = fitter(init_model, x, y, weights=np.zeros(10))


def test_linear_fit_model_set_weights():
    """Tests fitting multiple models simultaneously."""

    init_model = Polynomial1D(degree=2, c0=[1, 1], n_models=2)
    x = np.arange(10)
    y_expected = init_model(x, model_set_axis=False)
    assert y_expected.shape == (2, 10)

    # Add a bit of random noise
    with NumpyRNGContext(_RANDOM_SEED):
        y = y_expected + np.random.normal(0, 0.01, size=y_expected.shape)

    weights = np.ones_like(y)
    # Put a null weight for the min and max values
    weights[[0, 1], weights.argmin(axis=1)] = 0
    weights[[0, 1], weights.argmax(axis=1)] = 0
    fitter = LinearLSQFitter()
    fitted_model = fitter(init_model, x, y, weights=weights)
    assert_allclose(fitted_model(x, model_set_axis=False), y_expected,
                    rtol=1e-1)

    # Check that using null weights raises an error
    weights[0] = 0
    with pytest.raises(ValueError,
                       match='Found NaNs in the coefficient matrix'):
        with pytest.warns(RuntimeWarning,
                          match=r'invalid value encountered in.*divide'):
            fitted_model = fitter(init_model, x, y, weights=weights)

    # Now we mask the values where weight is 0
    with pytest.warns(RuntimeWarning,
                      match=r'invalid value encountered in.*divide'):
        fitted_model = fitter(init_model, x,
                              np.ma.array(y, mask=np.isclose(weights, 0)),
                              weights=weights)
    # Parameters for the first model are all NaNs
    assert np.all(np.isnan(fitted_model.param_sets[:, 0]))
    assert np.all(np.isnan(fitted_model(x, model_set_axis=False)[0]))
    # Second model is fitted correctly
    assert_allclose(fitted_model(x, model_set_axis=False)[1], y_expected[1],
                    rtol=1e-1)


def test_linear_fit_2d_model_set_errors():

    init_model = Polynomial2D(degree=2, c0_0=[1, 1], n_models=2)
    x = np.arange(10)
    y = np.arange(10)
    z = init_model(x, y, model_set_axis=False)

    fitter = LinearLSQFitter()
    with pytest.raises(ValueError):
        fitter(init_model, x[:5], y, z)
    with pytest.raises(ValueError):
        fitter(init_model, x, y, z[:, :5])


def test_linear_fit_2d_model_set_common_weight():
    init_model = Polynomial2D(degree=2, c1_0=[1, 2], c0_1=[-0.5, 1],
                              n_models=2,
                              fixed={'c1_0': True, 'c0_1': True})

    x, y = np.mgrid[0:5, 0:5]
    zz = np.array([1+x-0.5*y+0.1*x*x, 2*x+y-0.2*y*y])

    fitter = LinearLSQFitter()
    fitted_model = fitter(init_model, x, y, zz, weights=np.ones((5, 5)))

    assert_allclose(fitted_model(x, y, model_set_axis=False), zz,
                    atol=1e-14)


def test_linear_fit_flat_2d_model_set_common_weight():
    init_model = Polynomial2D(degree=2, c1_0=[1, 2], c0_1=[-0.5, 1],
                              n_models=2,
                              fixed={'c1_0': True, 'c0_1': True})

    x, y = np.mgrid[0:5, 0:5]
    x, y = x.flatten(), y.flatten()
    zz = np.array([1+x-0.5*y+0.1*x*x, 2*x+y-0.2*y*y])
    weights = np.ones(25)

    fitter = LinearLSQFitter()
    fitted_model = fitter(init_model, x, y, zz, weights=weights)

    assert_allclose(fitted_model(x, y, model_set_axis=False), zz,
                    atol=1e-14)


def test_linear_fit_2d_model_set_weights():
    init_model = Polynomial2D(degree=2, c1_0=[1, 2], c0_1=[-0.5, 1],
                              n_models=2,
                              fixed={'c1_0': True, 'c0_1': True})

    x, y = np.mgrid[0:5, 0:5]
    zz = np.array([1+x-0.5*y+0.1*x*x, 2*x+y-0.2*y*y])

    fitter = LinearLSQFitter()
    weights = [np.ones((5, 5)), np.ones((5, 5))]
    fitted_model = fitter(init_model, x, y, zz, weights=weights)

    assert_allclose(fitted_model(x, y, model_set_axis=False), zz,
                    atol=1e-14)


def test_linear_fit_flat_2d_model_set_weights():
    init_model = Polynomial2D(degree=2, c1_0=[1, 2], c0_1=[-0.5, 1],
                              n_models=2,
                              fixed={'c1_0': True, 'c0_1': True})

    x, y = np.mgrid[0:5, 0:5]
    x, y = x.flatten(), y.flatten()
    zz = np.array([1+x-0.5*y+0.1*x*x, 2*x+y-0.2*y*y])
    weights = np.ones((2, 25))

    fitter = LinearLSQFitter()
    fitted_model = fitter(init_model, x, y, zz, weights=weights)

    assert_allclose(fitted_model(x, y, model_set_axis=False), zz,
                    atol=1e-14)


class Test1ModelSet:
    """
    Check that fitting a single model works with a length-1 model set axis.
    It's not clear that this was originally intended usage, but it can be
    convenient, eg. when fitting a range of image rows that may be a single
    row, and some existing scripts might rely on it working.
    Currently this does not work with FittingWithOutlierRemoval.
    """

    def setup_class(self):
        self.x1 = np.arange(0, 10)
        self.y1 = np.array([0.5 + 2.5*self.x1])
        self.w1 = np.ones((10,))
        self.y1[0, 8] = 100.
        self.w1[8] = 0.
        self.y2, self.x2 = np.mgrid[0:10, 0:10]
        self.z2 = np.array([1 - 0.1*self.x2 + 0.2*self.y2])
        self.w2 = np.ones((10, 10))
        self.z2[0, 1, 2] = 100.
        self.w2[1, 2] = 0.

    def test_linear_1d_common_weights(self):
        model = Polynomial1D(1)
        fitter = LinearLSQFitter()
        model = fitter(model, self.x1, self.y1, weights=self.w1)
        assert_allclose(model.c0, 0.5, atol=1e-12)
        assert_allclose(model.c1, 2.5, atol=1e-12)

    def test_linear_1d_separate_weights(self):
        model = Polynomial1D(1)
        fitter = LinearLSQFitter()
        model = fitter(model, self.x1, self.y1,
                       weights=self.w1[np.newaxis, ...])
        assert_allclose(model.c0, 0.5, atol=1e-12)
        assert_allclose(model.c1, 2.5, atol=1e-12)

    def test_linear_1d_separate_weights_axis_1(self):
        model = Polynomial1D(1, model_set_axis=1)
        fitter = LinearLSQFitter()
        model = fitter(model, self.x1, self.y1.T,
                       weights=self.w1[..., np.newaxis])
        assert_allclose(model.c0, 0.5, atol=1e-12)
        assert_allclose(model.c1, 2.5, atol=1e-12)

    def test_linear_2d_common_weights(self):
        model = Polynomial2D(1)
        fitter = LinearLSQFitter()
        model = fitter(model, self.x2, self.y2, self.z2, weights=self.w2)
        assert_allclose(model.c0_0, 1., atol=1e-12)
        assert_allclose(model.c1_0, -0.1, atol=1e-12)
        assert_allclose(model.c0_1, 0.2, atol=1e-12)

    def test_linear_2d_separate_weights(self):
        model = Polynomial2D(1)
        fitter = LinearLSQFitter()
        model = fitter(model, self.x2, self.y2, self.z2,
                       weights=self.w2[np.newaxis, ...])
        assert_allclose(model.c0_0, 1., atol=1e-12)
        assert_allclose(model.c1_0, -0.1, atol=1e-12)
        assert_allclose(model.c0_1, 0.2, atol=1e-12)

    def test_linear_2d_separate_weights_axis_2(self):
        model = Polynomial2D(1, model_set_axis=2)
        fitter = LinearLSQFitter()
        model = fitter(model, self.x2, self.y2, np.rollaxis(self.z2, 0, 3),
                       weights=self.w2[..., np.newaxis])
        assert_allclose(model.c0_0, 1., atol=1e-12)
        assert_allclose(model.c1_0, -0.1, atol=1e-12)
        assert_allclose(model.c0_1, 0.2, atol=1e-12)