# Licensed under a 3-clause BSD style license - see LICENSE.rst
# pylint: disable=invalid-name, pointless-statement

import pickle

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

import astropy.units as u
from astropy.modeling.core import CompoundModel, Model, ModelDefinitionError
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.modeling.models import (
    Chebyshev1D, Chebyshev2D, Const1D, Gaussian1D, Gaussian2D, Identity, Legendre1D, Legendre2D,
    Linear1D, Mapping, Polynomial1D, Polynomial2D, Rotation2D, Scale, Shift, Tabular1D, fix_inputs)
from astropy.modeling.parameters import Parameter
from astropy.utils.compat.optional_deps import HAS_SCIPY  # noqa: F401


@pytest.mark.parametrize(('expr', 'result'),
                         [(lambda x, y: x + y, [5.0, 5.0]),
                          (lambda x, y: x - y, [-1.0, -1.0]),
                          (lambda x, y: x * y, [6.0, 6.0]),
                          (lambda x, y: x / y, [2.0 / 3.0, 2.0 / 3.0]),
                          (lambda x, y: x ** y, [8.0, 8.0])])
def test_model_set(expr, result):
    s = expr(Const1D((2, 2), n_models=2), Const1D((3, 3), n_models=2))
    out = s(0, model_set_axis=False)
    assert_array_equal(out, result)


@pytest.mark.parametrize(('expr', 'result'),
                         [(lambda x, y: x + y, [5.0, 5.0]),
                          (lambda x, y: x - y, [-1.0, -1.0]),
                          (lambda x, y: x * y, [6.0, 6.0]),
                          (lambda x, y: x / y, [2.0 / 3.0, 2.0 / 3.0]),
                          (lambda x, y: x ** y, [8.0, 8.0])])
def test_model_set_raises_value_error(expr, result):
    """Check that creating model sets with components whose _n_models are
       different raise a value error
    """
    with pytest.raises(ValueError):
        expr(Const1D((2, 2), n_models=2), Const1D(3, n_models=1))


@pytest.mark.parametrize(('expr', 'result'),
                         [(lambda x, y: x + y, 5.0),
                          (lambda x, y: x - y, -1.0),
                          (lambda x, y: x * y, 6.0),
                          (lambda x, y: x / y, 2.0 / 3.0),
                          (lambda x, y: x ** y, 8.0)])
def test_two_model_instance_arithmetic_1d(expr, result):
    """
    Like test_two_model_class_arithmetic_1d, but creates a new model from two
    model *instances* with fixed parameters.
    """

    s = expr(Const1D(2), Const1D(3))

    assert isinstance(s, CompoundModel)
    assert s.n_inputs == 1
    assert s.n_outputs == 1

    out = s(0)
    assert out == result

    assert isinstance(out, float)


def test_simple_two_model_compose_1d():
    """
    Shift and Scale are two of the simplest models to test model composition
    with.
    """

    S1 = Shift(2) | Scale(3)  # First shift then scale
    assert isinstance(S1, CompoundModel)
    assert S1.n_inputs == 1
    assert S1.n_outputs == 1
    assert S1(1) == 9.0

    S2 = Scale(2) | Shift(3)  # First scale then shift
    assert isinstance(S2, CompoundModel)
    assert S2.n_inputs == 1
    assert S2.n_outputs == 1
    assert S2(1) == 5.0

    # Test with array inputs
    assert_array_equal(S2([1, 2, 3]), [5.0, 7.0, 9.0])


def test_simple_two_model_compose_2d():
    """
    A simple example consisting of two rotations.
    """

    r1 = Rotation2D(45) | Rotation2D(45)

    assert isinstance(r1, CompoundModel)
    assert r1.n_inputs == 2
    assert r1.n_outputs == 2
    assert_allclose(r1(0, 1), (-1, 0), atol=1e-10)

    r2 = Rotation2D(90) | Rotation2D(90)  # Rotate twice by 90 degrees
    assert_allclose(r2(0, 1), (0, -1), atol=1e-10)

    # Compose R with itself to produce 4 rotations
    r3 = r1 | r1

    assert_allclose(r3(0, 1), (0, -1), atol=1e-10)


def test_n_submodels():
    """
    Test that CompoundModel.n_submodels properly returns the number
    of components.
    """

    g2 = Gaussian1D() + Gaussian1D()
    assert g2.n_submodels == 2

    g3 = g2 + Gaussian1D()
    assert g3.n_submodels == 3

    g5 = g3 | g2
    assert g5.n_submodels == 5

    g7 = g5 / g2
    assert g7.n_submodels == 7


def test_expression_formatting():
    """
    Test that the expression strings from compound models are formatted
    correctly.
    """

    # For the purposes of this test it doesn't matter a great deal what
    # model(s) are used in the expression, I don't think
    G = Gaussian1D(1, 1, 1)
    G2 = Gaussian2D(1, 2, 3, 4, 5, 6)

    M = G + G
    assert M._format_expression() == '[0] + [1]'

    M = G + G + G
    assert M._format_expression() == '[0] + [1] + [2]'

    M = G + G * G
    assert M._format_expression() == '[0] + [1] * [2]'

    M = G * G + G
    assert M._format_expression() == '[0] * [1] + [2]'

    M = G + G * G + G
    assert M._format_expression() == '[0] + [1] * [2] + [3]'

    M = (G + G) * (G + G)
    assert M._format_expression() == '([0] + [1]) * ([2] + [3])'

    # This example uses parentheses in the expression, but those won't be
    # preserved in the expression formatting since they technically aren't
    # necessary, and there's no way to know that they were originally
    # parenthesized (short of some deep, and probably not worthwhile
    # introspection)
    M = (G * G) + (G * G)
    assert M._format_expression() == '[0] * [1] + [2] * [3]'

    M = G ** G
    assert M._format_expression() == '[0] ** [1]'

    M = G + G ** G
    assert M._format_expression() == '[0] + [1] ** [2]'

    M = (G + G) ** G
    assert M._format_expression() == '([0] + [1]) ** [2]'

    M = G + G | G
    assert M._format_expression() == '[0] + [1] | [2]'

    M = G + (G | G)
    assert M._format_expression() == '[0] + ([1] | [2])'

    M = G & G | G2
    assert M._format_expression() == '[0] & [1] | [2]'

    M = G & (G | G)
    assert M._format_expression() == '[0] & ([1] | [2])'


def test_basic_compound_inverse():
    """
    Test basic inversion of compound models in the limited sense supported for
    models made from compositions and joins only.
    """

    t = (Shift(2) & Shift(3)) | (Scale(2) & Scale(3)) | Rotation2D(90)
    assert_allclose(t.inverse(*t(0, 1)), (0, 1))


@pytest.mark.parametrize('model', [
    Shift(0) + Shift(0) | Shift(0),
    Shift(0) - Shift(0) | Shift(0),
    Shift(0) * Shift(0) | Shift(0),
    Shift(0) / Shift(0) | Shift(0),
    Shift(0) ** Shift(0) | Shift(0),
    Gaussian1D(1, 2, 3) | Gaussian1D(4, 5, 6)])
def test_compound_unsupported_inverse(model):
    """
    Ensure inverses aren't supported in cases where it shouldn't be.
    """

    with pytest.raises(NotImplementedError):
        model.inverse


def test_mapping_basic_permutations():
    """
    Tests a couple basic examples of the Mapping model--specifically examples
    that merely permute the outputs.
    """

    x, y = Rotation2D(90)(1, 2)

    rs = Rotation2D(90) | Mapping((1, 0))
    x_prime, y_prime = rs(1, 2)
    assert_allclose((x, y), (y_prime, x_prime))

    # A more complicated permutation
    m = Rotation2D(90) & Scale(2)
    x, y, z = m(1, 2, 3)

    ms = m | Mapping((2, 0, 1))
    x_prime, y_prime, z_prime = ms(1, 2, 3)
    assert_allclose((x, y, z), (y_prime, z_prime, x_prime))


def test_mapping_inverse():
    """Tests inverting a compound model that includes a `Mapping`."""

    rs1 = Rotation2D(12.1) & Scale(13.2)
    rs2 = Rotation2D(14.3) & Scale(15.4)

    # Rotates 2 of the coordinates and scales the third--then rotates on a
    # different axis and scales on the axis of rotation.  No physical meaning
    # here just a simple test
    m = rs1 | Mapping([2, 0, 1]) | rs2

    assert_allclose((0, 1, 2), m.inverse(*m(0, 1, 2)), atol=1e-08)


def test_identity_input():
    """
    Test a case where an Identity (or Mapping) model is the first in a chain
    of composite models and thus is responsible for handling input broadcasting
    properly.

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

    ident1 = Identity(1)
    shift = Shift(1)
    rotation = Rotation2D(angle=90)
    model = ident1 & shift | rotation
    assert_allclose(model(1, 2), [-3.0, 1.0])


def test_invalid_operands():
    """
    Test that certain operators do not work with models whose inputs/outputs do
    not match up correctly.
    """

    with pytest.raises(ModelDefinitionError):
        Rotation2D(90) | Gaussian1D(1, 0, 0.1)

    with pytest.raises(ModelDefinitionError):
        Rotation2D(90) + Gaussian1D(1, 0, 0.1)


@pytest.mark.parametrize('poly', [Chebyshev2D(1, 2), Polynomial2D(2), Legendre2D(1, 2)])
def test_compound_with_polynomials_2d(poly):
    """
    Tests that polynomials are scaled when used in compound models.
    Issue #3699
    """
    poly.parameters = [1, 2, 3, 4, 1, 2]
    shift = Shift(3)
    model = poly | shift
    x, y = np.mgrid[:20, :37]
    result_compound = model(x, y)
    result = shift(poly(x, y))
    assert_allclose(result, result_compound)


def test_fix_inputs():
    g1 = Gaussian2D(1, 0, 0, 1, 2)
    g2 = Gaussian2D(1.5, .5, -.2, .5, .3)
    sg1_1 = fix_inputs(g1, {1: 0})
    assert_allclose(sg1_1(0), g1(0, 0))
    assert_allclose(sg1_1([0, 1, 3]), g1([0, 1, 3], [0, 0, 0]))
    sg1_2 = fix_inputs(g1, {'x': 1})
    assert_allclose(sg1_2(1.5), g1(1, 1.5))
    gg1 = g1 & g2
    sgg1_1 = fix_inputs(gg1, {1: 0.1, 3: 0.2})
    assert_allclose(sgg1_1(0, 0), gg1(0, 0.1, 0, 0.2))
    sgg1_2 = fix_inputs(gg1, {'x0': -.1, 2: .1})
    assert_allclose(sgg1_2(1, 1), gg1(-0.1, 1, 0.1, 1))
    assert_allclose(sgg1_2(y0=1, y1=1), gg1(-0.1, 1, 0.1, 1))


def test_fix_inputs_invalid():
    g1 = Gaussian2D(1, 0, 0, 1, 2)
    with pytest.raises(ValueError):
        fix_inputs(g1, {'x0': 0, 0: 0})

    with pytest.raises(ValueError):
        fix_inputs(g1, (0, 1))

    with pytest.raises(ValueError):
        fix_inputs(g1, {3: 2})

    with pytest.raises(ValueError):
        fix_inputs(g1, {np.int32(3): 2})

    with pytest.raises(ValueError):
        fix_inputs(g1, {np.int64(3): 2})

    with pytest.raises(ValueError):
        fix_inputs(g1, {'w': 2})

    with pytest.raises(ModelDefinitionError):
        CompoundModel('#', g1, g1)

    with pytest.raises(ValueError):
        gg1 = fix_inputs(g1, {0: 1})
        gg1(2, y=2)

    with pytest.raises(ValueError):
        gg1 = fix_inputs(g1, {np.int32(0): 1})
        gg1(2, y=2)

    with pytest.raises(ValueError):
        gg1 = fix_inputs(g1, {np.int64(0): 1})
        gg1(2, y=2)


def test_fix_inputs_with_bounding_box():
    g1 = Gaussian2D(1, 0, 0, 1, 1)
    g2 = Gaussian2D(1, 0, 0, 1, 1)
    assert g1.bounding_box == ((-5.5, 5.5), (-5.5, 5.5))

    gg1 = g1 & g2
    gg1.bounding_box = ((-5.5, 5.5), (-5.4, 5.4), (-5.3, 5.3), (-5.2, 5.2))
    assert gg1.bounding_box == ((-5.5, 5.5), (-5.4, 5.4), (-5.3, 5.3), (-5.2, 5.2))

    sg = fix_inputs(gg1, {0: 0, 2: 0})
    assert sg.bounding_box == ((-5.5, 5.5), (-5.3, 5.3))

    g1 = Gaussian1D(10, 3, 1)
    g = g1 & g1
    g.bounding_box = ((1, 4), (6, 8))
    gf = fix_inputs(g, {0: 1})
    assert gf.bounding_box == (1, 4)


def test_indexing_on_instance():
    """Test indexing on compound model instances."""

    m = Gaussian1D(1, 0, 0.1) + Const1D(2)
    assert isinstance(m[0], Gaussian1D)
    assert isinstance(m[1], Const1D)
    assert m.param_names == ('amplitude_0', 'mean_0', 'stddev_0', 'amplitude_1')

    # Test parameter equivalence
    assert m[0].amplitude == 1 == m.amplitude_0
    assert m[0].mean == 0 == m.mean_0
    assert m[0].stddev == 0.1 == m.stddev_0
    assert m[1].amplitude == 2 == m.amplitude_1

    # Test that parameter value updates are symmetric between the compound
    # model and the submodel returned by indexing
    const = m[1]
    m.amplitude_1 = 42
    assert const.amplitude == 42
    const.amplitude = 137
    assert m.amplitude_1 == 137

    # Similar couple of tests, but now where the compound model was created
    # from model instances
    g = Gaussian1D(1, 2, 3, name='g')
    p = Polynomial1D(2, name='p')
    m = g + p
    assert m[0].name == 'g'
    assert m[1].name == 'p'
    assert m['g'].name == 'g'
    assert m['p'].name == 'p'

    poly = m[1]
    m.c0_1 = 12345
    assert poly.c0 == 12345
    poly.c1 = 6789
    assert m.c1_1 == 6789

    # Test negative indexing
    assert isinstance(m[-1], Polynomial1D)
    assert isinstance(m[-2], Gaussian1D)

    with pytest.raises(IndexError):
        m[42]

    with pytest.raises(IndexError):
        m['foobar']

    # Confirm index-by-name works with fix_inputs
    g = Gaussian2D(1, 2, 3, 4, 5, name='g')
    m = fix_inputs(g, {0: 1})
    assert m['g'].name == 'g'

    # Test string slicing
    A = Const1D(1.1, name='A')
    B = Const1D(2.1, name='B')
    C = Const1D(3.1, name='C')
    M = A + B * C
    assert_allclose(M['B':'C'](1), 6.510000000000001)


class _ConstraintsTestA(Model):
    stddev = Parameter(default=0, min=0, max=0.3)
    mean = Parameter(default=0, fixed=True)

    @staticmethod
    def evaluate(stddev, mean):
        return stddev, mean


class _ConstraintsTestB(Model):
    mean = Parameter(default=0, fixed=True)

    @staticmethod
    def evaluate(mean):
        return mean


def test_inherit_constraints():
    """
    Various tests for copying of constraint values between compound models and
    their members.

    Regression test for https://github.com/astropy/astropy/issues/3481
    """
    model = (Gaussian1D(bounds={'stddev': (0, 0.3)}, fixed={'mean': True}) +
             Gaussian1D(fixed={'mean': True}))

    # Lots of assertions in this test as there are multiple interfaces to
    # parameter constraints

    assert 'stddev_0' in model.bounds
    assert model.bounds['stddev_0'] == (0, 0.3)
    assert model.stddev_0.bounds == (0, 0.3)
    assert 'mean_0' in model.fixed
    assert model.fixed['mean_0'] is True
    assert model.mean_0.fixed is True
    assert 'mean_1' in model.fixed
    assert model.fixed['mean_1'] is True
    assert model.mean_1.fixed is True

    assert model.stddev_0 is model[0].stddev
    # Great, all the constraints were inherited properly
    # Now what about if we update them through the sub-models?
    model.stddev_0.bounds = (0, 0.4)
    assert model[0].stddev.bounds == (0, 0.4)
    assert model[0].bounds['stddev'] == (0, 0.4)

    model.stddev_0.bounds = (0.1, 0.5)
    assert model[0].stddev.bounds == (0.1, 0.5)
    assert model[0].bounds['stddev'] == (0.1, 0.5)

    model[1].mean.fixed = False
    assert model.mean_1.fixed is False
    assert model[1].mean.fixed is False

    # Now turn off syncing of constraints
    assert model.bounds['stddev_0'] == (0.1, 0.5)
    model.sync_constraints = False
    model[0].stddev.bounds = (0, 0.2)
    assert model.bounds['stddev_0'] == (0.1, 0.5)
    model.sync_constraints = True
    assert model.bounds['stddev_0'] == (0, 0.2)


def test_compound_custom_inverse():
    """
    Test that a compound model with a custom inverse has that inverse applied
    when the inverse of another model, of which it is a component, is computed.
    Regression test for https://github.com/astropy/astropy/issues/3542
    """

    poly = Polynomial1D(1, c0=1, c1=2)
    scale = Scale(1)
    shift = Shift(1)

    model1 = poly | scale
    model1.inverse = poly

    # model1 now has a custom inverse (the polynomial itself, ignoring the
    # trivial scale factor)
    model2 = shift | model1

    assert_allclose(model2.inverse(1), (poly | shift.inverse)(1))

    # Make sure an inverse is not allowed if the models were combined with the
    # wrong operator, or if one of the models doesn't have an inverse defined
    with pytest.raises(NotImplementedError):
        (shift + model1).inverse

    with pytest.raises(NotImplementedError):
        (model1 & poly).inverse


def test_pickle_compound():
    """
    Regression test for
    https://github.com/astropy/astropy/issues/3867#issuecomment-114547228
    """

    # Test pickling a compound model instance
    g1 = Gaussian1D(1.0, 0.0, 0.1)
    g2 = Gaussian1D([2.0, 3.0], [0.0, 0.0], [0.2, 0.3])
    m = g1 + g2
    m2 = pickle.loads(pickle.dumps(m))
    assert m.param_names == m2.param_names
    assert m.__class__.__name__ == m2.__class__.__name__
    assert np.all(m.parameters == m2.parameters)
    assert np.all(m(0) == m2(0))


def test_update_parameters():
    offx = Shift(1)
    scl = Scale(2)
    m = offx | scl
    assert m(1) == 4

    offx.offset = 42
    assert m(1) == 86

    m.factor_1 = 100
    assert m(1) == 4300
    m2 = m | offx
    assert m2(1) == 4342


def test_name():
    offx = Shift(1)
    scl = Scale(2)
    m = offx | scl
    scl.name = "scale"
    assert m.submodel_names == ('None_0', 'scale')
    assert m.name is None
    m.name = "M"
    assert m.name == "M"
    m1 = m.rename("M1")
    assert m.name == "M1"
    assert m1.name == "M1"


def test_name_index():
    g1 = Gaussian1D(1, 1, 1)
    g2 = Gaussian1D(1, 2, 1)
    g = g1 + g2
    with pytest.raises(IndexError):
        g['bozo']
    g1.name = 'bozo'
    assert g['bozo'].mean == 1
    g2.name = 'bozo'
    with pytest.raises(IndexError):
        g['bozo']


@pytest.mark.skipif("not HAS_SCIPY")
def test_tabular_in_compound():
    """
    Issue #7411 - evaluate should not change the shape of the output.
    """
    t = Tabular1D(points=([1, 5, 7],), lookup_table=[12, 15, 19],
                  bounds_error=False)
    rot = Rotation2D(2)
    p = Polynomial1D(1)
    x = np.arange(12).reshape((3, 4))
    # Create a compound model which does not execute Tabular.__call__,
    # but model.evaluate and is followed by a Rotation2D which
    # checks the exact shapes.
    model = p & t | rot
    x1, y1 = model(x, x)
    assert x1.ndim == 2
    assert y1.ndim == 2


def test_bounding_box():
    g = Gaussian2D() + Gaussian2D(2, .5, .1, 2, 3, 0)
    g.bounding_box = ((0, 1), (0, .5))
    y, x = np.mgrid[0:10, 0:10]
    y = y / 3.
    x = x / 3.
    val = g(x, y, with_bounding_box=True)
    compare = np.array([
        [2.93738984, 2.93792011, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [2.87857153, 2.88188761, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [2.70492922, 2.71529265, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [2.45969972, 2.47912103, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan,
         np.nan, np.nan, np.nan, np.nan, np.nan]])
    mask = ~np.isnan(val)
    assert_allclose(val[mask], compare[mask])
    val2 = g(x+2, y+2, with_bounding_box=True)
    assert np.isnan(val2).sum() == 100
    # val3 = g(.1, .1, with_bounding_box=True)


@pytest.mark.skipif("not HAS_SCIPY")
def test_bounding_box_with_units():
    points = np.arange(5) * u.pix
    lt = np.arange(5) * u.AA
    t = Tabular1D(points, lt)

    assert t(1 * u.pix, with_bounding_box=True) == 1. * u.AA


@pytest.mark.parametrize('poly', [Chebyshev1D(5), Legendre1D(5), Polynomial1D(5)])
def test_compound_with_polynomials_1d(poly):
    """
    Tests that polynomials are offset when used in compound models.
    Issue #3699
    """
    poly.parameters = [1, 2, 3, 4, 1, 2]
    shift = Shift(3)
    model = poly | shift
    x = np.linspace(-5, 5, 10)
    result_compound = model(x)
    result = shift(poly(x))
    assert_allclose(result, result_compound)
    assert model.param_names == ('c0_0', 'c1_0', 'c2_0', 'c3_0', 'c4_0', 'c5_0', 'offset_1')


def test_replace_submodel():
    """
    Replace a model in a Compound model
    """
    S1 = Shift(2, name='shift2') | Scale(3, name='scale3')  # First shift then scale
    S2 = Scale(2, name='scale2') | Shift(3, name='shift3')  # First scale then shift

    m = S1 & S2
    assert m(1, 2) == (9, 7)

    m2 = m.replace_submodel('scale3', Scale(4, name='scale4'))
    assert m2(1, 2) == (12, 7)
    assert m(1, 2) == (9, 7)
    # Check the inverse has been updated
    assert m2.inverse(12, 7) == (1, 2)

    # Produce the same result by replacing a single model with a compound
    m3 = m.replace_submodel('shift2', Shift(2) | Scale(2))
    assert m(1, 2) == (9, 7)
    assert m3(1, 2) == (18, 7)
    # Check the inverse has been updated
    assert m3.inverse(18, 7) == (1, 2)

    # Test with arithmetic model compunding operator
    m = S1 + S2
    assert m(1) == 14
    m2 = m.replace_submodel('scale2', Scale(4, name='scale4'))
    assert m2(1) == 16

    # Test with fix_inputs()
    R = fix_inputs(Rotation2D(angle=90, name='rotate'), {0: 1})
    m4 = S1 | R
    assert_allclose(m4(0), (-6, 1))

    m5 = m4.replace_submodel('rotate', Rotation2D(180))
    assert_allclose(m5(0), (-1, -6))

    # Check we get a value error when model name doesn't exist
    with pytest.raises(ValueError):
        m2 = m.replace_submodel('not_there', Scale(2))

    # And now a model set
    P = Polynomial1D(degree=1, n_models=2, name='poly')
    S = Shift([1, 2], n_models=2)
    m = P | S
    assert_array_equal(m([0, 1]), (1, 2))
    with pytest.raises(ValueError):
        m2 = m.replace_submodel('poly', Polynomial1D(degree=1, c0=1))
    m2 = m.replace_submodel('poly', Polynomial1D(degree=1, c0=[1, 2],
                                                 n_models=2))
    assert_array_equal(m2([0, 1]), (2, 4))

    # Ensure previous _user_inverse doesn't stick around
    S1 = Shift(1)
    S2 = Shift(2)
    S3 = Shift(3, name='S3')

    S23 = S2 | S3
    S23.inverse = Shift(-4.9)
    m = S1 & S23

    # This should delete the S23._user_inverse
    m2 = m.replace_submodel('S3', Shift(4))
    assert m2(1, 2) == (2, 8)
    assert m2.inverse(2, 8) == (1, 2)


@pytest.mark.parametrize(
    "expr",
    [
        lambda m1, m2: m1 + m2,
        lambda m1, m2: m1 - m2,
        lambda m1, m2: m1 * m2,
        lambda m1, m2: m1 / m2,
    ],
)
def test_compound_evaluate(expr):
    """
    Tests that compound evaluate function produces the same
    result as the models with the operator applied
    """
    x = np.linspace(-5, 5, 10)
    # Some evaluate functions assume that inputs are numpy arrays or quantities including Const1D
    p1 = np.array([1, 2, 3, 4, 1, 2])
    p2 = np.array([1, 0, 0.5])

    model1 = Polynomial1D(5)
    model2 = Gaussian1D(2, 1, 5)
    compound = expr(model1, model2)

    assert_array_equal(
        compound.evaluate(x, *p1, *p2),
        expr(model1.evaluate(x, *p1), model2.evaluate(x, *p2)),
    )


def test_compound_evaluate_power():
    """
    Tests that compound evaluate function produces the same
    result as the models with the power operator applied
    """
    x = np.linspace(-5, 5, 10)
    p1 = np.array([1, 0, 0.2])
    p2 = np.array([3])

    model1 = Gaussian1D(2, 1, 5)
    model2 = Const1D(2)
    compound = model1 ** model2

    assert_array_equal(
        compound.evaluate(x, *p1, *p2),
        model1.evaluate(x, *p1) ** model2.evaluate(x, *p2),
    )


def test_compound_evaluate_double_shift():
    x = np.linspace(-5, 5, 10)
    y = np.linspace(-5, 5, 10)

    m1 = Gaussian2D(1, 0, 0, 1, 1, 1)
    m2 = Shift(1)
    m3 = Shift(2)
    m = Gaussian2D(1, 0, 0, 1, 1, 1) & Shift(1) & Shift(2)
    assert_array_equal(
        m.evaluate(x, y, x - 10, y + 20, 1, 0, 0, 1, 1, 1, 1, 2),
        [
            m1.evaluate(x, y, 1, 0, 0, 1, 1, 1),
            m2.evaluate(x - 10, 1),
            m3.evaluate(y + 20, 2),
        ],
    )


@pytest.mark.parametrize(
    "expr",
    [
        lambda m1, m2: m1 + m2,
        lambda m1, m2: m1 - m2,
        lambda m1, m2: m1 * m2,
        lambda m1, m2: m1 / m2,
    ],
)
def test_compound_evaluate_named_param(expr):
    """
    Tests that compound evaluate function produces the same
    result as the models with the operator applied
    """
    x = np.linspace(-5, 5, 10)
    p1 = np.array([1, 0, 0.2])
    p2 = np.array([3, 0.5, 0.5])

    model1 = Gaussian1D(2, 1, 5)
    model2 = Gaussian1D(2, 1, 5)
    compound = expr(model1, model2)

    assert_array_equal(
        compound.evaluate(
            x, *p2, amplitude_0=p1[0], mean_0=p1[1], stddev_0=p1[2]
        ),
        expr(model1.evaluate(x, *p1), model2.evaluate(x, *p2)),
    )


def test_compound_evaluate_name_param_power():
    """
    Tests that compound evaluate function produces the same
    result as the models with the power operator applied
    """
    x = np.linspace(-5, 5, 10)
    p1 = np.array([1, 0, 0.2])
    p2 = np.array([3])

    model1 = Gaussian1D(2, 1, 5)
    model2 = Const1D(2)
    compound = model1 ** model2

    assert_array_equal(
        compound.evaluate(
            x, *p2, amplitude_0=p1[0], mean_0=p1[1], stddev_0=p1[2]
        ),
        model1.evaluate(x, *p1) ** model2.evaluate(x, *p2),
    )


def test_compound_evaluate_and():
    """
    Tests that compound evaluate function produces the same
    result as the models with the operator applied
    """
    x = np.linspace(-5, 5, 10)
    p1 = np.array([1, 0.1, 0.5])
    p2 = np.array([3])

    model1 = Gaussian1D()
    model2 = Shift()
    compound = model1 & model2

    assert_array_equal(
        compound.evaluate(x, x, *p1, p2),
        [model1.evaluate(x, *p1), model2.evaluate(x, p2)],
    )


def test_compound_evaluate_or():
    """
    Tests that compound evaluate function produces the same
    result as the models with the operator applied
    """
    x = np.linspace(-5, 5, 10)
    p1 = np.array([0.5])
    p2_amplitude = np.array([3])
    p2_mean = np.array([0])
    p2_std = np.array([0.1])

    model1 = Shift(0.5)
    model2 = Gaussian1D(1, 0, 0.5)
    compound = model1 | model2

    assert_array_equal(
        compound.evaluate(x, p1, p2_amplitude, p2_mean, p2_std),
        model2.evaluate(model1.evaluate(x, p1), p2_amplitude, p2_mean, p2_std),
    )


def test_compound_evaluate_fix_inputs_by_keyword():
    """
    Tests that compound evaluate function produces the same
    result as the models fix_inputs operator is applied
    when using the keyword
    """
    y, x = np.mgrid[:10, :10]

    model_params = [3, 0, 0.1, 1, 0.5, 0]

    model = Gaussian2D(1, 2, 0, 0.5)
    compound = fix_inputs(model, {"x": x + 5})

    assert_array_equal(
        compound.evaluate(x, y, *model_params),
        model.evaluate(x + 5, y, *model_params),
    )


def test_compound_evaluate_fix_inputs_by_position():
    """
    Tests that compound evaluate function produces the same
    result as the models fix_inputs operator is applied
    when using the input index
    """
    y, x = np.mgrid[:10, :10]

    model_params = [3, 0, 0.1, 1, 0.5, 0]

    model = Gaussian2D(1, 2, 0, 0.5)
    compound = fix_inputs(model, {0: x + 5})

    assert_array_equal(
        compound.evaluate(x, y, *model_params),
        model.evaluate(x + 5, y, *model_params),
    )


@pytest.mark.skipif('not HAS_SCIPY')
def test_fit_multiplied_compound_model_with_mixed_units():
    """
    Regression test for issue #12320
    """

    fitter = LevMarLSQFitter()
    x = np.linspace(0, 1, 101) * u.s
    y = np.linspace(5, 10, 101) * u.m * u.kg / u.s

    m1 = Linear1D(slope=5*u.m/u.s/u.s, intercept=1.0*u.m/u.s)
    m2 = Linear1D(slope=0.0*u.kg/u.s, intercept=10.0*u.kg)
    truth = m1 * m2
    fit = fitter(truth, x, y)

    unfit_output = truth(x)
    fit_output = fit(x)

    assert unfit_output.unit == fit_output.unit == (u.kg * u.m / u.s)
    assert_allclose(unfit_output, fit_output)

    for name in truth.param_names:
        assert getattr(truth, name) == getattr(fit, name)


@pytest.mark.skipif('not HAS_SCIPY')
def test_fit_multiplied_recursive_compound_model_with_mixed_units():
    """
    Regression test for issue #12320
    """

    fitter = LevMarLSQFitter()

    x = np.linspace(0, 1, 101) * u.s
    y = np.linspace(5, 10, 101) * u.m * u.m * u.kg / u.s

    m1 = Linear1D(slope=5*u.m/u.s/u.s, intercept=1.0*u.m/u.s)
    m2 = Linear1D(slope=0.0*u.kg/u.s, intercept=10.0*u.kg)
    m3 = Linear1D(slope=0.0*u.m/u.s, intercept=10.0*u.m)
    truth = m1 * m2 * m3
    fit = fitter(truth, x, y)

    unfit_output = truth(x)
    fit_output = fit(x)

    assert unfit_output.unit == fit_output.unit == (u.kg * u.m * u.m / u.s)
    assert_allclose(unfit_output, fit_output)

    for name in truth.param_names:
        assert getattr(truth, name) == getattr(fit, name)

    x = np.linspace(0, 1, 101) * u.s
    y = np.linspace(5, 10, 101) * u.m * u.m * u.kg * u.kg / u.s

    m1 = Linear1D(slope=5*u.m/u.s/u.s, intercept=1.0*u.m/u.s)
    m2 = Linear1D(slope=0.0*u.kg/u.s, intercept=10.0*u.kg)
    m3 = Linear1D(slope=0.0*u.m/u.s, intercept=10.0*u.m)
    m4 = Linear1D(slope=0.0*u.kg/u.s, intercept=10.0*u.kg)
    m11 = m1 * m2
    m22 = m3 * m4
    truth = m11 * m22
    fit = fitter(truth, x, y)

    unfit_output = truth(x)
    fit_output = fit(x)

    assert unfit_output.unit == fit_output.unit == (u.kg * u.kg * u.m * u.m / u.s)
    assert_allclose(unfit_output, fit_output)

    for name in truth.param_names:
        assert getattr(truth, name) == getattr(fit, name)


@pytest.mark.skipif('not HAS_SCIPY')
def test_fit_divided_compound_model_with_mixed_units():
    """
    Regression test for issue #12320
    """

    fitter = LevMarLSQFitter()
    x = np.linspace(0, 1, 101) * u.s
    y = np.linspace(5, 10, 101) * u.kg * u.m / u.s

    m1 = Linear1D(slope=5*u.kg*u.m/u.s, intercept=1.0*u.kg*u.m)
    m2 = Linear1D(slope=0.0*u.s/u.s, intercept=10.0*u.s)
    truth = m1 / m2
    fit = fitter(truth, x, y)

    unfit_output = truth(x)
    fit_output = fit(x)

    assert unfit_output.unit == fit_output.unit == (u.kg * u.m / u.s)
    assert_allclose(unfit_output, fit_output)

    for name in truth.param_names:
        assert getattr(truth, name) == getattr(fit, name)


@pytest.mark.skipif('not HAS_SCIPY')
def test_fit_mixed_recursive_compound_model_with_mixed_units():
    """
    Regression test for issue #12320
    """

    fitter = LevMarLSQFitter()

    x = np.linspace(0, 1, 101) * u.s
    y = np.linspace(5, 10, 101) * u.kg * u.m * u.m / u.s

    m1 = Linear1D(slope=5*u.kg*u.m/u.s, intercept=1.0*u.kg*u.m)
    m2 = Linear1D(slope=0.0*u.s/u.s, intercept=10.0*u.s)
    m3 = Linear1D(slope=0.0*u.m/u.s, intercept=10.0*u.m)
    truth = m1 / m2 * m3
    fit = fitter(truth, x, y)

    unfit_output = truth(x)
    fit_output = fit(x)

    assert unfit_output.unit == fit_output.unit == (u.kg * u.m * u.m / u.s)
    assert_allclose(unfit_output, fit_output)

    for name in truth.param_names:
        assert getattr(truth, name) == getattr(fit, name)

    x = np.linspace(0, 1, 101) * u.s
    y = np.linspace(5, 10, 101) * u.kg * u.kg * u.m * u.m / u.s

    m1 = Linear1D(slope=5*u.kg*u.m/u.s, intercept=1.0*u.kg*u.m)
    m2 = Linear1D(slope=0.0*u.s/u.s, intercept=10.0*u.s)
    m3 = Linear1D(slope=0.0*u.m/u.s, intercept=10.0*u.m)
    m4 = Linear1D(slope=0.0*u.kg/u.s, intercept=10.0*u.kg)
    m11 = m1 / m2
    m22 = m3 * m4
    truth = m11 * m22
    fit = fitter(truth, x, y)

    unfit_output = truth(x)
    fit_output = fit(x)

    assert unfit_output.unit == fit_output.unit == (u.kg * u.kg * u.m * u.m / u.s)
    assert_allclose(unfit_output, fit_output)

    for name in truth.param_names:
        assert getattr(truth, name) == getattr(fit, name)
