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

import pickle
import pytest

import numpy as np

from numpy.testing import assert_allclose, assert_array_equal

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


@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