# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
This module tests fitting and model evaluation with various inputs
"""

import numpy as np
import pytest
from numpy.testing import assert_allclose

from astropy.modeling import fitting, models
from astropy.modeling.core import Fittable1DModel, FittableModel, Model
from astropy.modeling.parameters import Parameter
from astropy.utils.compat.optional_deps import HAS_SCIPY

model1d_params = [
    (models.Polynomial1D, [2]),
    (models.Legendre1D, [2]),
    (models.Chebyshev1D, [2]),
    (models.Shift, [2]),
    (models.Scale, [2]),
]

model2d_params = [
    (models.Polynomial2D, [2]),
    (models.Legendre2D, [1, 2]),
    (models.Chebyshev2D, [1, 2]),
]

fitters = [
    fitting.LevMarLSQFitter,
    fitting.TRFLSQFitter,
    fitting.LMLSQFitter,
    fitting.DogBoxLSQFitter,
]


class TestInputType:
    """
    This class tests that models accept numbers, lists and arrays.

    Add new models to one of the lists above to test for this.
    """

    def setup_class(self):
        self.x = 5.3
        self.y = 6.7
        self.x1 = np.arange(1, 10, 0.1)
        self.y1 = np.arange(1, 10, 0.1)
        self.y2, self.x2 = np.mgrid[:10, :8]

    @pytest.mark.parametrize(("model", "params"), model1d_params)
    def test_input1D(self, model, params):
        m = model(*params)
        m(self.x)
        m(self.x1)
        m(self.x2)

    @pytest.mark.parametrize(("model", "params"), model2d_params)
    def test_input2D(self, model, params):
        m = model(*params)
        m(self.x, self.y)
        m(self.x1, self.y1)
        m(self.x2, self.y2)


class TestFitting:
    """Test various input options to fitting routines."""

    def setup_class(self):
        self.x1 = np.arange(10)
        self.y, self.x = np.mgrid[:10, :10]

    def test_linear_fitter_1set(self):
        """1 set 1D x, 1pset"""

        expected = np.array([0, 1, 1, 1])
        p1 = models.Polynomial1D(3)
        p1.parameters = [0, 1, 1, 1]
        y1 = p1(self.x1)
        pfit = fitting.LinearLSQFitter()
        model = pfit(p1, self.x1, y1)
        assert_allclose(model.parameters, expected, atol=10 ** (-7))

    def test_linear_fitter_Nset(self):
        """1 set 1D x, 2 sets 1D y, 2 param_sets"""

        expected = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
        p1 = models.Polynomial1D(3, n_models=2)
        p1.parameters = [0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0]
        params = {}
        for i in range(4):
            params[p1.param_names[i]] = [i, i]
        p1 = models.Polynomial1D(3, model_set_axis=0, **params)
        y1 = p1(self.x1, model_set_axis=False)
        pfit = fitting.LinearLSQFitter()
        model = pfit(p1, self.x1, y1)
        assert_allclose(model.param_sets, expected, atol=10 ** (-7))

    def test_linear_fitter_1dcheb(self):
        """1 pset, 1 set 1D x, 1 set 1D y, Chebyshev 1D polynomial"""

        expected = np.array(
            [
                [
                    2817.2499999999995,
                    4226.6249999999991,
                    1680.7500000000009,
                    273.37499999999926,
                ]
            ]
        ).T
        ch1 = models.Chebyshev1D(3)
        ch1.parameters = [0, 1, 2, 3]
        y1 = ch1(self.x1)
        pfit = fitting.LinearLSQFitter()
        model = pfit(ch1, self.x1, y1)
        assert_allclose(model.param_sets, expected, atol=10 ** (-2))

    def test_linear_fitter_1dlegend(self):
        """
        1 pset, 1 set 1D x, 1 set 1D y, Legendre 1D polynomial
        """

        expected = np.array(
            [
                [
                    1925.5000000000011,
                    3444.7500000000005,
                    1883.2500000000014,
                    364.4999999999996,
                ]
            ]
        ).T
        leg1 = models.Legendre1D(3)
        leg1.parameters = [1, 2, 3, 4]
        y1 = leg1(self.x1)
        pfit = fitting.LinearLSQFitter()
        model = pfit(leg1, self.x1, y1)
        assert_allclose(model.param_sets, expected, atol=10 ** (-12))

    def test_linear_fitter_1set2d(self):
        p2 = models.Polynomial2D(2)
        p2.parameters = [0, 1, 2, 3, 4, 5]
        expected = [0, 1, 2, 3, 4, 5]
        z = p2(self.x, self.y)
        pfit = fitting.LinearLSQFitter()
        model = pfit(p2, self.x, self.y, z)
        assert_allclose(model.parameters, expected, atol=10 ** (-12))
        assert_allclose(model(self.x, self.y), z, atol=10 ** (-12))

    def test_wrong_numpset(self):
        """
        A ValueError is raised if a 1 data set (1d x, 1d y) is fit
        with a model with multiple parameter sets.
        """

        MESSAGE = (
            r"Number of data sets .* is expected to equal the number of parameter sets"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            p1 = models.Polynomial1D(5)
            y1 = p1(self.x1)
            p1 = models.Polynomial1D(5, n_models=2)
            pfit = fitting.LinearLSQFitter()
            pfit(p1, self.x1, y1)

    def test_wrong_pset(self):
        """A case of 1 set of x and multiple sets of y and parameters."""

        expected = np.array(
            [
                [1, 0],
                [1, 1],
                [1, 2],
                [1, 3],
                [1, 4],
                [1, 5],
            ]
        )
        p1 = models.Polynomial1D(5, n_models=2)
        params = {}
        for i in range(6):
            params[p1.param_names[i]] = [1, i]
        p1 = models.Polynomial1D(5, model_set_axis=0, **params)
        y1 = p1(self.x1, model_set_axis=False)
        pfit = fitting.LinearLSQFitter()
        model = pfit(p1, self.x1, y1)
        assert_allclose(model.param_sets, expected, atol=10 ** (-7))

    @pytest.mark.skipif(not HAS_SCIPY, reason="requires scipy")
    @pytest.mark.parametrize("fitter", fitters)
    def test_nonlinear_lsqt_1set_1d(self, fitter):
        """1 set 1D x, 1 set 1D y, 1 pset NonLinearFitter"""
        fitter = fitter()

        g1 = models.Gaussian1D(10, mean=3, stddev=0.2)
        y1 = g1(self.x1)
        model = fitter(g1, self.x1, y1)
        assert_allclose(model.parameters, [10, 3, 0.2])

    @pytest.mark.skipif(not HAS_SCIPY, reason="requires scipy")
    @pytest.mark.parametrize("fitter", fitters)
    def test_nonlinear_lsqt_Nset_1d(self, fitter):
        """1 set 1D x, 1 set 1D y, 2 param_sets, NonLinearFitter"""
        fitter = fitter()

        MESSAGE = r"Non-linear fitters can only fit one data set at a time"
        with pytest.raises(ValueError, match=MESSAGE):
            g1 = models.Gaussian1D(
                [10.2, 10], mean=[3, 3.2], stddev=[0.23, 0.2], n_models=2
            )
            y1 = g1(self.x1, model_set_axis=False)
            _ = fitter(g1, self.x1, y1)

    @pytest.mark.skipif(not HAS_SCIPY, reason="requires scipy")
    @pytest.mark.parametrize("fitter", fitters)
    def test_nonlinear_lsqt_1set_2d(self, fitter):
        """1 set 2d x, 1set 2D y, 1 pset, NonLinearFitter"""
        fitter = fitter()

        g2 = models.Gaussian2D(
            10, x_mean=3, y_mean=4, x_stddev=0.3, y_stddev=0.2, theta=0
        )
        z = g2(self.x, self.y)
        model = fitter(g2, self.x, self.y, z)
        assert_allclose(model.parameters, [10, 3, 4, 0.3, 0.2, 0])

    @pytest.mark.skipif(not HAS_SCIPY, reason="requires scipy")
    @pytest.mark.parametrize("fitter", fitters)
    def test_nonlinear_lsqt_Nset_2d(self, fitter):
        """1 set 2d x, 1set 2D y, 2 param_sets, NonLinearFitter"""
        fitter = fitter()

        MESSAGE = (
            r"Input argument .* does not have the correct dimensions in .* for a model"
            r" set with .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            g2 = models.Gaussian2D(
                [10, 10],
                [3, 3],
                [4, 4],
                x_stddev=[0.3, 0.3],
                y_stddev=[0.2, 0.2],
                theta=[0, 0],
                n_models=2,
            )
            z = g2(self.x.flatten(), self.y.flatten())
            _ = fitter(g2, self.x, self.y, z)


class TestEvaluation:
    """
    Test various input options to model evaluation

    TestFitting actually covers evaluation of polynomials
    """

    def setup_class(self):
        self.x1 = np.arange(20)
        self.y, self.x = np.mgrid[:10, :10]

    def test_non_linear_NYset(self):
        """
        This case covers:
            N param sets , 1 set 1D x --> N 1D y data
        """

        g1 = models.Gaussian1D([10, 10], [3, 3], [0.2, 0.2], n_models=2)
        y1 = g1(self.x1, model_set_axis=False)
        assert np.all((y1[0, :] - y1[1, :]).nonzero() == np.array([]))

    def test_non_linear_NXYset(self):
        """
        This case covers: N param sets , N sets 1D x --> N N sets 1D y data
        """

        g1 = models.Gaussian1D([10, 10], [3, 3], [0.2, 0.2], n_models=2)
        xx = np.array([self.x1, self.x1])
        y1 = g1(xx)
        assert_allclose(y1[:, 0], y1[:, 1], atol=10 ** (-12))

    def test_p1_1set_1pset(self):
        """1 data set, 1 pset, Polynomial1D"""

        p1 = models.Polynomial1D(4)
        y1 = p1(self.x1)
        assert y1.shape == (20,)

    def test_p1_nset_npset(self):
        """N data sets, N param_sets, Polynomial1D"""

        p1 = models.Polynomial1D(4, n_models=2)
        y1 = p1(np.array([self.x1, self.x1]).T, model_set_axis=-1)
        assert y1.shape == (20, 2)
        assert_allclose(y1[0, :], y1[1, :], atol=10 ** (-12))

    def test_p2_1set_1pset(self):
        """1 pset, 1 2D data set, Polynomial2D"""

        p2 = models.Polynomial2D(5)
        z = p2(self.x, self.y)
        assert z.shape == (10, 10)

    def test_p2_nset_npset(self):
        """N param_sets, N 2D data sets, Poly2d"""

        p2 = models.Polynomial2D(5, n_models=2)
        xx = np.array([self.x, self.x])
        yy = np.array([self.y, self.y])
        z = p2(xx, yy)
        assert z.shape == (2, 10, 10)

    def test_nset_domain(self):
        """
        Test model set with negative model_set_axis.

        In this case model_set_axis=-1 is identical to model_set_axis=1.
        """

        xx = np.array([self.x1, self.x1]).T
        xx[0, 0] = 100
        xx[1, 0] = 100
        xx[2, 0] = 99
        p1 = models.Polynomial1D(5, c0=[1, 2], c1=[3, 4], n_models=2)
        yy = p1(xx, model_set_axis=-1)
        assert_allclose(xx.shape, yy.shape)
        yy1 = p1(xx, model_set_axis=1)
        assert_allclose(yy, yy1)

    def test_evaluate_gauss2d(self):
        cov = np.array([[1.0, 0.8], [0.8, 3]])
        g = models.Gaussian2D(1.0, 5.0, 4.0, cov_matrix=cov)
        y, x = np.mgrid[:10, :10]
        g(x, y)


class TModel_1_1(Fittable1DModel):
    p1 = Parameter()
    p2 = Parameter()

    @staticmethod
    def evaluate(x, p1, p2):
        return x + p1 + p2


class TestSingleInputSingleOutputSingleModel:
    """
    A suite of tests to check various cases of parameter and input combinations
    on models with n_input = n_output = 1 on a toy model with n_models=1.

    Many of these tests mirror test cases in
    ``astropy.modeling.tests.test_parameters.TestParameterInitialization``,
    except that this tests how different parameter arrangements interact with
    different types of model inputs.
    """

    def test_scalar_parameters_scalar_input(self):
        """
        Scalar parameters with a scalar input should return a scalar.
        """

        t = TModel_1_1(1, 10)
        y = t(100)
        assert isinstance(y, float)
        assert np.ndim(y) == 0
        assert y == 111

    def test_scalar_parameters_1d_array_input(self):
        """
        Scalar parameters should broadcast with an array input to result in an
        array output of the same shape as the input.
        """

        t = TModel_1_1(1, 10)
        y = t(np.arange(5) * 100)
        assert isinstance(y, np.ndarray)
        assert np.shape(y) == (5,)
        assert np.all(y == [11, 111, 211, 311, 411])

    def test_scalar_parameters_2d_array_input(self):
        """
        Scalar parameters should broadcast with an array input to result in an
        array output of the same shape as the input.
        """

        t = TModel_1_1(1, 10)
        y = t(np.arange(6).reshape(2, 3) * 100)
        assert isinstance(y, np.ndarray)
        assert np.shape(y) == (2, 3)
        assert np.all(y == [[11, 111, 211], [311, 411, 511]])

    def test_scalar_parameters_3d_array_input(self):
        """
        Scalar parameters should broadcast with an array input to result in an
        array output of the same shape as the input.
        """

        t = TModel_1_1(1, 10)
        y = t(np.arange(12).reshape(2, 3, 2) * 100)
        assert isinstance(y, np.ndarray)
        assert np.shape(y) == (2, 3, 2)
        assert np.all(
            y
            == [
                [[11, 111], [211, 311], [411, 511]],
                [[611, 711], [811, 911], [1011, 1111]],
            ]
        )

    def test_1d_array_parameters_scalar_input(self):
        """
        Array parameters should all be broadcastable with each other, and with
        a scalar input the output should be broadcast to the maximum dimensions
        of the parameters.
        """

        t = TModel_1_1([1, 2], [10, 20])
        y = t(100)
        assert isinstance(y, np.ndarray)
        assert np.shape(y) == (2,)
        assert np.all(y == [111, 122])

    def test_1d_array_parameters_1d_array_input(self):
        """
        When given an array input it must be broadcastable with all the
        parameters.
        """

        t = TModel_1_1([1, 2], [10, 20])
        y1 = t([100, 200])
        assert np.shape(y1) == (2,)
        assert np.all(y1 == [111, 222])

        y2 = t([[100], [200]])
        assert np.shape(y2) == (2, 2)
        assert np.all(y2 == [[111, 122], [211, 222]])

        MESSAGE = (
            r"self input argument 'x' of shape .* cannot be broadcast with parameter"
            r" 'p1' of shape .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            # Doesn't broadcast
            t([100, 200, 300])

    def test_2d_array_parameters_2d_array_input(self):
        """
        When given an array input it must be broadcastable with all the
        parameters.
        """

        t = TModel_1_1([[1, 2], [3, 4]], [[10, 20], [30, 40]])

        y1 = t([[100, 200], [300, 400]])
        assert np.shape(y1) == (2, 2)
        assert np.all(y1 == [[111, 222], [333, 444]])

        y2 = t([[[[100]], [[200]]], [[[300]], [[400]]]])
        assert np.shape(y2) == (2, 2, 2, 2)
        assert np.all(
            y2
            == [
                [[[111, 122], [133, 144]], [[211, 222], [233, 244]]],
                [[[311, 322], [333, 344]], [[411, 422], [433, 444]]],
            ]
        )

        MESSAGE = (
            r"self input argument .* of shape .* cannot be broadcast with parameter .*"
            r" of shape .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            # Doesn't broadcast
            t([[100, 200, 300], [400, 500, 600]])

    def test_mixed_array_parameters_1d_array_input(self):
        """
        When given an array input it must be broadcastable with all the
        parameters.
        """

        t = TModel_1_1(
            [
                [[0.01, 0.02, 0.03], [0.04, 0.05, 0.06]],
                [[0.07, 0.08, 0.09], [0.10, 0.11, 0.12]],
            ],
            [1, 2, 3],
        )

        y1 = t([10, 20, 30])
        assert np.shape(y1) == (2, 2, 3)
        assert_allclose(
            y1,
            [
                [[11.01, 22.02, 33.03], [11.04, 22.05, 33.06]],
                [[11.07, 22.08, 33.09], [11.10, 22.11, 33.12]],
            ],
        )

        y2 = t([[[[10]]], [[[20]]], [[[30]]]])
        assert np.shape(y2) == (3, 2, 2, 3)
        assert_allclose(
            y2,
            [
                [
                    [[11.01, 12.02, 13.03], [11.04, 12.05, 13.06]],
                    [[11.07, 12.08, 13.09], [11.10, 12.11, 13.12]],
                ],
                [
                    [[21.01, 22.02, 23.03], [21.04, 22.05, 23.06]],
                    [[21.07, 22.08, 23.09], [21.10, 22.11, 23.12]],
                ],
                [
                    [[31.01, 32.02, 33.03], [31.04, 32.05, 33.06]],
                    [[31.07, 32.08, 33.09], [31.10, 32.11, 33.12]],
                ],
            ],
        )


class TestSingleInputSingleOutputTwoModel:
    """
    A suite of tests to check various cases of parameter and input combinations
    on models with n_input = n_output = 1 on a toy model with n_models=2.

    Many of these tests mirror test cases in
    ``astropy.modeling.tests.test_parameters.TestParameterInitialization``,
    except that this tests how different parameter arrangements interact with
    different types of model inputs.

    With n_models=2 all outputs should have a first dimension of size 2 (unless
    defined with model_set_axis != 0).
    """

    def test_scalar_parameters_scalar_input(self):
        """
        Scalar parameters with a scalar input should return a 1-D array with
        size equal to the number of models.
        """

        t = TModel_1_1([1, 2], [10, 20], n_models=2)

        y = t(100)
        assert np.shape(y) == (2,)
        assert np.all(y == [111, 122])

    def test_scalar_parameters_1d_array_input(self):
        """
        The dimension of the input should match the number of models unless
        model_set_axis=False is given, in which case the input is copied across
        all models.
        """

        t = TModel_1_1([1, 2], [10, 20], n_models=2)

        MESSAGE = (
            r"Input argument .* does not have the correct dimensions in .* for a model"
            r" set with .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            t(np.arange(5) * 100)

        y1 = t([100, 200])
        assert np.shape(y1) == (2,)
        assert np.all(y1 == [111, 222])

        y2 = t([100, 200], model_set_axis=False)
        # In this case the value [100, 200, 300] should be evaluated on each
        # model rather than evaluating the first model with 100 and the second
        # model  with 200
        assert np.shape(y2) == (2, 2)
        assert np.all(y2 == [[111, 211], [122, 222]])

        y3 = t([100, 200, 300], model_set_axis=False)
        assert np.shape(y3) == (2, 3)
        assert np.all(y3 == [[111, 211, 311], [122, 222, 322]])

    def test_scalar_parameters_2d_array_input(self):
        """
        The dimension of the input should match the number of models unless
        model_set_axis=False is given, in which case the input is copied across
        all models.
        """

        t = TModel_1_1([1, 2], [10, 20], n_models=2)

        y1 = t(np.arange(6).reshape(2, 3) * 100)
        assert np.shape(y1) == (2, 3)
        assert np.all(y1 == [[11, 111, 211], [322, 422, 522]])

        y2 = t(np.arange(6).reshape(2, 3) * 100, model_set_axis=False)
        assert np.shape(y2) == (2, 2, 3)
        assert np.all(
            y2 == [[[11, 111, 211], [311, 411, 511]], [[22, 122, 222], [322, 422, 522]]]
        )

    def test_scalar_parameters_3d_array_input(self):
        """
        The dimension of the input should match the number of models unless
        model_set_axis=False is given, in which case the input is copied across
        all models.
        """

        t = TModel_1_1([1, 2], [10, 20], n_models=2)
        data = np.arange(12).reshape(2, 3, 2) * 100

        y1 = t(data)
        assert np.shape(y1) == (2, 3, 2)
        assert np.all(
            y1
            == [
                [[11, 111], [211, 311], [411, 511]],
                [[622, 722], [822, 922], [1022, 1122]],
            ]
        )

        y2 = t(data, model_set_axis=False)
        assert np.shape(y2) == (2, 2, 3, 2)
        assert np.all(y2 == np.array([data + 11, data + 22]))

    def test_1d_array_parameters_scalar_input(self):
        """
        Array parameters should all be broadcastable with each other, and with
        a scalar input the output should be broadcast to the maximum dimensions
        of the parameters.
        """

        t = TModel_1_1([[1, 2, 3], [4, 5, 6]], [[10, 20, 30], [40, 50, 60]], n_models=2)

        y = t(100)
        assert np.shape(y) == (2, 3)
        assert np.all(y == [[111, 122, 133], [144, 155, 166]])

    def test_1d_array_parameters_1d_array_input(self):
        """
        When the input is an array, if model_set_axis=False then it must
        broadcast with the shapes of the parameters (excluding the
        model_set_axis).

        Otherwise all dimensions must be broadcastable.
        """

        t = TModel_1_1([[1, 2, 3], [4, 5, 6]], [[10, 20, 30], [40, 50, 60]], n_models=2)

        MESSAGE = (
            r"Input argument .* does not have the correct dimensions in .* for a model"
            r" set with .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            y1 = t([100, 200, 300])

        y1 = t([100, 200])
        assert np.shape(y1) == (2, 3)
        assert np.all(y1 == [[111, 122, 133], [244, 255, 266]])

        MESSAGE = (
            r"Model input argument .* of shape .* cannot be broadcast with parameter .*"
            r" of shape .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            # Doesn't broadcast with the shape of the parameters, (3,)
            y2 = t([100, 200], model_set_axis=False)

        y2 = t([100, 200, 300], model_set_axis=False)
        assert np.shape(y2) == (2, 3)
        assert np.all(y2 == [[111, 222, 333], [144, 255, 366]])

    def test_2d_array_parameters_2d_array_input(self):
        t = TModel_1_1(
            [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
            [[[10, 20], [30, 40]], [[50, 60], [70, 80]]],
            n_models=2,
        )
        y1 = t([[100, 200], [300, 400]])
        assert np.shape(y1) == (2, 2, 2)
        assert np.all(
            y1
            == [
                [[111, 222], [133, 244]],
                [[355, 466], [377, 488]],
            ]
        )

        MESSAGE = (
            r"Model input argument .* of shape .* cannot be broadcast with parameter .*"
            r" of shape .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            y2 = t([[100, 200, 300], [400, 500, 600]])

        y2 = t([[[100, 200], [300, 400]], [[500, 600], [700, 800]]])
        assert np.shape(y2) == (2, 2, 2)
        assert np.all(
            y2
            == [
                [[111, 222], [333, 444]],
                [[555, 666], [777, 888]],
            ]
        )

    def test_mixed_array_parameters_1d_array_input(self):
        t = TModel_1_1(
            [
                [[0.01, 0.02, 0.03], [0.04, 0.05, 0.06]],
                [[0.07, 0.08, 0.09], [0.10, 0.11, 0.12]],
            ],
            [[1, 2, 3], [4, 5, 6]],
            n_models=2,
        )

        MESSAGE = (
            r"Input argument .* does not have the correct dimensions in .* for a model"
            r" set with .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            y = t([10, 20, 30])

        y = t([10, 20, 30], model_set_axis=False)
        assert np.shape(y) == (2, 2, 3)
        assert_allclose(
            y,
            [
                [[11.01, 22.02, 33.03], [11.04, 22.05, 33.06]],
                [[14.07, 25.08, 36.09], [14.10, 25.11, 36.12]],
            ],
        )


class TModel_1_2(FittableModel):
    n_inputs = 1
    n_outputs = 2

    p1 = Parameter()
    p2 = Parameter()
    p3 = Parameter()

    @staticmethod
    def evaluate(x, p1, p2, p3):
        return (x + p1 + p2, x + p1 + p2 + p3)


class TestSingleInputDoubleOutputSingleModel:
    """
    A suite of tests to check various cases of parameter and input combinations
    on models with n_input = 1 but n_output = 2 on a toy model with n_models=1.

    As of writing there are not enough controls to adjust how outputs from such
    a model should be formatted (currently the shapes of outputs are assumed to
    be directly associated with the shapes of corresponding inputs when
    n_inputs == n_outputs).  For now, the approach taken for cases like this is
    to assume all outputs should have the same format.
    """

    def test_scalar_parameters_scalar_input(self):
        """
        Scalar parameters with a scalar input should return a scalar.
        """

        t = TModel_1_2(1, 10, 1000)
        y, z = t(100)
        assert isinstance(y, float)
        assert isinstance(z, float)
        assert np.ndim(y) == np.ndim(z) == 0
        assert y == 111
        assert z == 1111

    def test_scalar_parameters_1d_array_input(self):
        """
        Scalar parameters should broadcast with an array input to result in an
        array output of the same shape as the input.
        """

        t = TModel_1_2(1, 10, 1000)
        y, z = t(np.arange(5) * 100)
        assert isinstance(y, np.ndarray)
        assert isinstance(z, np.ndarray)
        assert np.shape(y) == np.shape(z) == (5,)
        assert np.all(y == [11, 111, 211, 311, 411])
        assert np.all(z == (y + 1000))

    def test_scalar_parameters_2d_array_input(self):
        """
        Scalar parameters should broadcast with an array input to result in an
        array output of the same shape as the input.
        """

        t = TModel_1_2(1, 10, 1000)
        y, z = t(np.arange(6).reshape(2, 3) * 100)
        assert isinstance(y, np.ndarray)
        assert isinstance(z, np.ndarray)
        assert np.shape(y) == np.shape(z) == (2, 3)
        assert np.all(y == [[11, 111, 211], [311, 411, 511]])
        assert np.all(z == (y + 1000))

    def test_scalar_parameters_3d_array_input(self):
        """
        Scalar parameters should broadcast with an array input to result in an
        array output of the same shape as the input.
        """

        t = TModel_1_2(1, 10, 1000)
        y, z = t(np.arange(12).reshape(2, 3, 2) * 100)
        assert isinstance(y, np.ndarray)
        assert isinstance(z, np.ndarray)
        assert np.shape(y) == np.shape(z) == (2, 3, 2)
        assert np.all(
            y
            == [
                [[11, 111], [211, 311], [411, 511]],
                [[611, 711], [811, 911], [1011, 1111]],
            ]
        )
        assert np.all(z == (y + 1000))

    def test_1d_array_parameters_scalar_input(self):
        """
        Array parameters should all be broadcastable with each other, and with
        a scalar input the output should be broadcast to the maximum dimensions
        of the parameters.
        """

        t = TModel_1_2([1, 2], [10, 20], [1000, 2000])
        y, z = t(100)
        assert isinstance(y, np.ndarray)
        assert isinstance(z, np.ndarray)
        assert np.shape(y) == np.shape(z) == (2,)
        assert np.all(y == [111, 122])
        assert np.all(z == [1111, 2122])

    def test_1d_array_parameters_1d_array_input(self):
        """
        When given an array input it must be broadcastable with all the
        parameters.
        """

        t = TModel_1_2([1, 2], [10, 20], [1000, 2000])
        y1, z1 = t([100, 200])
        assert np.shape(y1) == np.shape(z1) == (2,)
        assert np.all(y1 == [111, 222])
        assert np.all(z1 == [1111, 2222])

        y2, z2 = t([[100], [200]])
        assert np.shape(y2) == np.shape(z2) == (2, 2)
        assert np.all(y2 == [[111, 122], [211, 222]])
        assert np.all(z2 == [[1111, 2122], [1211, 2222]])

        MESSAGE = (
            r"self input argument .* of shape .* cannot be broadcast with parameter .*"
            r" of shape .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            # Doesn't broadcast
            y3, z3 = t([100, 200, 300])

    def test_2d_array_parameters_2d_array_input(self):
        """
        When given an array input it must be broadcastable with all the
        parameters.
        """

        t = TModel_1_2(
            [[1, 2], [3, 4]], [[10, 20], [30, 40]], [[1000, 2000], [3000, 4000]]
        )

        y1, z1 = t([[100, 200], [300, 400]])
        assert np.shape(y1) == np.shape(z1) == (2, 2)
        assert np.all(y1 == [[111, 222], [333, 444]])
        assert np.all(z1 == [[1111, 2222], [3333, 4444]])

        y2, z2 = t([[[[100]], [[200]]], [[[300]], [[400]]]])
        assert np.shape(y2) == np.shape(z2) == (2, 2, 2, 2)
        assert np.all(
            y2
            == [
                [[[111, 122], [133, 144]], [[211, 222], [233, 244]]],
                [[[311, 322], [333, 344]], [[411, 422], [433, 444]]],
            ]
        )
        assert np.all(
            z2
            == [
                [[[1111, 2122], [3133, 4144]], [[1211, 2222], [3233, 4244]]],
                [[[1311, 2322], [3333, 4344]], [[1411, 2422], [3433, 4444]]],
            ]
        )

        MESSAGE = (
            r"self input argument .* of shape .* cannot be broadcast with parameter .*"
            r" of shape .*"
        )
        with pytest.raises(ValueError, match=MESSAGE):
            # Doesn't broadcast
            y3, z3 = t([[100, 200, 300], [400, 500, 600]])

    def test_mixed_array_parameters_1d_array_input(self):
        """
        When given an array input it must be broadcastable with all the
        parameters.
        """

        t = TModel_1_2(
            [
                [[0.01, 0.02, 0.03], [0.04, 0.05, 0.06]],
                [[0.07, 0.08, 0.09], [0.10, 0.11, 0.12]],
            ],
            [1, 2, 3],
            [100, 200, 300],
        )

        y1, z1 = t([10, 20, 30])
        assert np.shape(y1) == np.shape(z1) == (2, 2, 3)
        assert_allclose(
            y1,
            [
                [[11.01, 22.02, 33.03], [11.04, 22.05, 33.06]],
                [[11.07, 22.08, 33.09], [11.10, 22.11, 33.12]],
            ],
        )
        assert_allclose(
            z1,
            [
                [[111.01, 222.02, 333.03], [111.04, 222.05, 333.06]],
                [[111.07, 222.08, 333.09], [111.10, 222.11, 333.12]],
            ],
        )

        y2, z2 = t([[[[10]]], [[[20]]], [[[30]]]])
        assert np.shape(y2) == np.shape(z2) == (3, 2, 2, 3)
        assert_allclose(
            y2,
            [
                [
                    [[11.01, 12.02, 13.03], [11.04, 12.05, 13.06]],
                    [[11.07, 12.08, 13.09], [11.10, 12.11, 13.12]],
                ],
                [
                    [[21.01, 22.02, 23.03], [21.04, 22.05, 23.06]],
                    [[21.07, 22.08, 23.09], [21.10, 22.11, 23.12]],
                ],
                [
                    [[31.01, 32.02, 33.03], [31.04, 32.05, 33.06]],
                    [[31.07, 32.08, 33.09], [31.10, 32.11, 33.12]],
                ],
            ],
        )

        assert_allclose(
            z2,
            [
                [
                    [[111.01, 212.02, 313.03], [111.04, 212.05, 313.06]],
                    [[111.07, 212.08, 313.09], [111.10, 212.11, 313.12]],
                ],
                [
                    [[121.01, 222.02, 323.03], [121.04, 222.05, 323.06]],
                    [[121.07, 222.08, 323.09], [121.10, 222.11, 323.12]],
                ],
                [
                    [[131.01, 232.02, 333.03], [131.04, 232.05, 333.06]],
                    [[131.07, 232.08, 333.09], [131.10, 232.11, 333.12]],
                ],
            ],
        )


# test broadcasting rules
broadcast_models = [
    {"model": models.Identity(2), "inputs": [0, [1, 1]], "outputs": [0, [1, 1]]},
    {"model": models.Identity(2), "inputs": [[1, 1], 0], "outputs": [[1, 1], 0]},
    {"model": models.Mapping((0, 1)), "inputs": [0, [1, 1]], "outputs": [0, [1, 1]]},
    {"model": models.Mapping((1, 0)), "inputs": [0, [1, 1]], "outputs": [[1, 1], 0]},
    {
        "model": models.Mapping((1, 0), n_inputs=3),
        "inputs": [0, [1, 1], 2],
        "outputs": [[1, 1], 0],
    },
    {
        "model": models.Mapping((0, 1, 0)),
        "inputs": [0, [1, 1]],
        "outputs": [0, [1, 1], 0],
    },
    {
        "model": models.Mapping((0, 1, 1)),
        "inputs": [0, [1, 1]],
        "outputs": [0, [1, 1], [1, 1]],
    },
    {"model": models.Polynomial2D(1, c0_0=1), "inputs": [0, [1, 1]], "outputs": [1, 1]},
    {"model": models.Polynomial2D(1, c0_0=1), "inputs": [0, 1], "outputs": 1},
    {
        "model": models.Gaussian2D(1, 1, 2, 1, 1.2),
        "inputs": [0, [1, 1]],
        "outputs": [0.42860385, 0.42860385],
    },
    {
        "model": models.Gaussian2D(1, 1, 2, 1, 1.2),
        "inputs": [0, 1],
        "outputs": 0.428603846153,
    },
    {
        "model": models.Polynomial2D(1, c0_0=1) & models.Polynomial2D(1, c0_0=2),
        "inputs": [1, 1, 1, 1],
        "outputs": (1, 2),
    },
    {
        "model": models.Polynomial2D(1, c0_0=1) & models.Polynomial2D(1, c0_0=2),
        "inputs": [1, 1, [1, 1], [1, 1]],
        "outputs": (1, [2, 2]),
    },
    {
        "model": models.math.MultiplyUfunc(),
        "inputs": [np.array([np.linspace(0, 1, 5)]).T, np.arange(2)],
        "outputs": np.array(
            [[0.0, 0.0], [0.0, 0.25], [0.0, 0.5], [0.0, 0.75], [0.0, 1.0]]
        ),
    },
]


@pytest.mark.parametrize("model", broadcast_models)
def test_mixed_input(model):
    result = model["model"](*model["inputs"])
    if np.isscalar(result):
        assert_allclose(result, model["outputs"])
    else:
        for i in range(len(result)):
            assert_allclose(result[i], model["outputs"][i])


def test_more_outputs():
    class M(FittableModel):
        standard_broadcasting = False
        n_inputs = 2
        n_outputs = 3

        a = Parameter()

        def evaluate(self, x, y, a):
            return a * x, a - x, a + y

        def __call__(self, *args, **kwargs):
            inputs, _ = super().prepare_inputs(*args, **kwargs)

            outputs = self.evaluate(*inputs, *self.parameters)
            output_shapes = [out.shape for out in outputs]
            output_shapes = [() if shape == (1,) else shape for shape in output_shapes]
            return self.prepare_outputs((tuple(output_shapes),), *outputs, **kwargs)

    c = M(1)
    result = c([1, 1], 1)
    expected = [[1.0, 1.0], [0.0, 0.0], 2.0]
    for r, e in zip(result, expected):
        assert_allclose(r, e)

    c = M(1)
    result = c(1, [1, 1])
    expected = [1.0, 0.0, [2.0, 2.0]]
    for r, e in zip(result, expected):
        assert_allclose(r, e)


class TInputFormatter(Model):
    """
    A toy model to test input/output formatting.
    """

    n_inputs = 2
    n_outputs = 2
    outputs = ("x", "y")

    @staticmethod
    def evaluate(x, y):
        return x, y


def test_format_input_scalars():
    model = TInputFormatter()
    result = model(1, 2)
    assert result == (1, 2)


def test_format_input_arrays():
    model = TInputFormatter()
    result = model([1, 1], [2, 2])
    assert_allclose(result, (np.array([1, 1]), np.array([2, 2])))


def test_format_input_arrays_transposed():
    model = TInputFormatter()
    input = np.array([[1, 1]]).T, np.array([[2, 2]]).T
    result = model(*input)
    assert_allclose(result, input)


@pytest.mark.parametrize(
    "model",
    [
        models.Gaussian2D(),
        models.Polynomial2D(1),
        models.Rotation2D(),
        models.Pix2Sky_TAN(),
        models.Tabular2D(lookup_table=np.ones((4, 5))),
    ],
)
@pytest.mark.skipif(not HAS_SCIPY, reason="requires scipy")
def test_call_keyword_args_1(model):
    """
    Test calling a model with positional, keywrd and a mixture of both arguments.
    """
    positional = model(1, 2)
    assert_allclose(positional, model(x=1, y=2))
    assert_allclose(positional, model(1, y=2))

    model.inputs = ("r", "t")
    assert_allclose(positional, model(r=1, t=2))
    assert_allclose(positional, model(1, t=2))
    assert_allclose(positional, model(1, 2))

    MESSAGE = r"Too many input arguments - expected 2, got .*"
    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, 3)

    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, t=12, r=3)

    MESSAGE = r"Missing input arguments - expected 2, got 1"
    with pytest.raises(ValueError, match=MESSAGE):
        model(1)


@pytest.mark.parametrize(
    "model",
    [
        models.Gaussian1D(),
        models.Polynomial1D(1),
        models.Tabular1D(lookup_table=np.ones((5,))),
    ],
)
@pytest.mark.skipif(not HAS_SCIPY, reason="requires scipy")
def test_call_keyword_args_2(model):
    """
    Test calling a model with positional, keywrd and a mixture of both arguments.
    """
    positional = model(1)
    assert_allclose(positional, model(x=1))

    model.inputs = ("r",)
    assert_allclose(positional, model(r=1))

    MESSAGE = r"Too many input arguments - expected .*, got .*"
    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, 3)

    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, t=12, r=3)

    MESSAGE = r"Missing input arguments - expected 1, got 0"
    with pytest.raises(ValueError, match=MESSAGE):
        model()


@pytest.mark.parametrize(
    "model",
    [
        models.Gaussian2D() | models.Polynomial1D(1),
        models.Gaussian1D() & models.Polynomial1D(1),
        models.Gaussian2D() + models.Polynomial2D(1),
        models.Gaussian2D() - models.Polynomial2D(1),
        models.Gaussian2D() * models.Polynomial2D(1),
        models.Identity(2) | models.Polynomial2D(1),
        models.Mapping((1,)) | models.Polynomial1D(1),
    ],
)
def test_call_keyword_args_3(model):
    """
    Test calling a model with positional, keywrd and a mixture of both arguments.
    """
    positional = model(1, 2)
    model.inputs = ("r", "t")
    assert_allclose(positional, model(r=1, t=2))

    assert_allclose(positional, model(1, t=2))

    MESSAGE = r"Too many input arguments - expected .*, got .*"
    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, 3)

    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, t=12, r=3)

    MESSAGE = r"Missing input arguments - expected 2, got 0"
    with pytest.raises(ValueError, match=MESSAGE):
        model()


@pytest.mark.parametrize(
    "model",
    [
        models.Identity(2),
        models.Mapping((0, 1)),
        models.Mapping((1,)),
    ],
)
def test_call_keyword_mappings(model):
    """
    Test calling a model with positional, keywrd and a mixture of both arguments.
    """
    positional = model(1, 2)
    assert_allclose(positional, model(x0=1, x1=2))
    assert_allclose(positional, model(1, x1=2))

    # We take a copy before modifying the model since otherwise this changes
    # the instance used in the parametrize call and affects future test runs.
    model = model.copy()

    model.inputs = ("r", "t")
    assert_allclose(positional, model(r=1, t=2))
    assert_allclose(positional, model(1, t=2))
    assert_allclose(positional, model(1, 2))

    MESSAGE = r"Too many input arguments - expected .*, got .*"
    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, 3)

    with pytest.raises(ValueError, match=MESSAGE):
        model(1, 2, t=12, r=3)

    MESSAGE = r"Missing input arguments - expected 2, got 1"
    with pytest.raises(ValueError, match=MESSAGE):
        model(1)
