import copy
import os
import pickle
import time

import numpy as np
import pytest

import aesara
import aesara.gpuarray
import aesara.tensor as aet
from aesara.compile.function import function
from aesara.compile.function.types import UnusedInputError
from aesara.compile.io import In, Out
from aesara.compile.mode import Mode
from aesara.configdefaults import config
from aesara.graph.basic import Constant
from aesara.graph.fg import MissingInputError
from aesara.graph.opt import OpKeyOptimizer, PatternSub
from aesara.tensor.math import dot
from aesara.tensor.math import sum as aet_sum
from aesara.tensor.math import tanh
from aesara.tensor.type import (
    dmatrix,
    dscalar,
    dscalars,
    dvector,
    fmatrix,
    fscalar,
    iscalar,
    matrix,
    scalar,
    scalars,
    vector,
)
from aesara.utils import exc_message


def PatternOptimizer(p1, p2, ign=True):
    return OpKeyOptimizer(PatternSub(p1, p2), ignore_newtrees=ign)


class TestFunction:
    @pytest.mark.xfail()
    def test_none(self):
        fn = function([], None)  # ok
        rval = fn()
        assert (
            rval != []
        ), "See #254: Using None as function output leads to [] return value"
        assert rval is None

    def test_empty(self):
        fn = function([], [])  # ok
        assert fn() == []

    def test_extra_inputs(self):
        x, s = scalars("xs")
        fn = function([x], [x])
        with pytest.raises(TypeError):
            fn(1, 2)

    def test_missing_inputs(self):
        def fn():
            x, s = scalars("xs")
            function([], [x])

        with pytest.raises(MissingInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            # Ignore unused input s, as it hides the other error
            function([s], [x], on_unused_input="ignore")

        with pytest.raises(MissingInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            function([s], [x])

        with pytest.raises(UnusedInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            # Ignore unused input s, as it hides the other error
            function([s], x, on_unused_input="ignore")

        with pytest.raises(MissingInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            function([s], x)

        with pytest.raises(UnusedInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            # Ignore unused input s, as it hides the other error
            function([s], Out(x), on_unused_input="ignore")

        with pytest.raises(MissingInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            function([s], Out(x))

        with pytest.raises(UnusedInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            function([In(x, update=s + x)], x)

        with pytest.raises(MissingInputError):
            fn()

        def fn():
            x, s = scalars("xs")
            function([In(x, update=((s * s) + x))], x)

        with pytest.raises(MissingInputError):
            fn()

    def test_input_anon_singleton(self):
        x, s = scalars("xs")
        fn = function([s, x], [x + s])
        assert fn(2, 3) == [5]
        # no state
        assert fn(2, 3) == [5]

    def test_input_anon_unpack(self):
        x, s = scalars("xs")
        fn = function([s, x], x + s)
        assert fn(2, 3) == 5

    def test_naming_rule0(self):
        x, s = scalars("xs")
        f = function([x, s], x / s)
        assert f(1, 2) == 0.5
        assert f(2, 1) == 2.0
        assert f(s=2, x=1) == 0.5
        assert f(x=2, s=1) == 2.0
        assert f(2, s=1) == 2.0

        with pytest.raises(TypeError):
            # got multiple values for keyword argument 'x'
            f(2, x=2.0)
        with pytest.raises(TypeError):
            # takes exactly 2 non-keyword arguments (1 given)
            f(x=1)
        with pytest.raises(TypeError):
            # takes exactly 2 non-keyword arguments (0 given)
            f(s=1)

    def test_naming_rule1(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")
        f = function([a, s], a / s)
        assert f(1, 2) == 0.5
        assert f(2, 1) == 2.0
        assert f(2, s=1) == 2.0

        with pytest.raises(TypeError):
            # got unexpected keyword argument 'q'
            f(q=2, s=1)
        with pytest.raises(TypeError):
            # got unexpected keyword argument 'a'
            f(a=2, s=1)

    def test_naming_rule2(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        # x's name is ignored because it is followed by anonymous parameter a.
        # Ignore unused input x, as it hides the other error
        f = function([x, a, s], a / s, on_unused_input="ignore")
        assert f(9, 1, 2) == 0.5
        assert f(9, 2, 1) == 2.0
        assert f(9, 2, s=1) == 2.0

        with pytest.raises(TypeError):
            # got unexpected keyword argument 'x'
            f(x=9, a=2, s=1)
        with pytest.raises(TypeError):
            # got unexpected keyword argument 'x'
            f(5.0, x=9)

    def test_naming_rule3(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        # x's name is not ignored (as in test_naming_rule2) because a has a default value.
        f = function([x, In(a, value=1.0), s], a / s + x)
        assert f(9, 2, 4) == 9.5  # can specify all args in order
        assert f(9, 2, s=4) == 9.5  # can give s as kwarg
        assert f(9, s=4) == 9.25  # can give s as kwarg, get default a
        assert f(x=9, s=4) == 9.25  # can give s as kwarg, omit a, x as kw
        with pytest.raises(TypeError):
            # got unexpected keyword argument 'a'
            f(x=9, a=2, s=4)
        with pytest.raises(TypeError):
            # takes exactly 3 non-keyword arguments (0 given)
            f()
        with pytest.raises(TypeError):
            # takes exactly 3 non-keyword arguments (1 given)
            f(x=9)

    def test_naming_rule4(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function([x, In(a, value=1.0, name="a"), s], a / s + x)

        assert f(9, 2, 4) == 9.5  # can specify all args in order
        assert f(9, 2, s=4) == 9.5  # can give s as kwarg
        assert f(9, s=4) == 9.25  # can give s as kwarg, get default a
        assert f(9, a=2, s=4) == 9.5  # can give s as kwarg, a as kwarg
        assert f(x=9, a=2, s=4) == 9.5  # can give all kwargs
        assert f(x=9, s=4) == 9.25  # can give all kwargs
        with pytest.raises(TypeError):
            # takes exactly 3 non-keyword arguments (0 given)
            f()
        with pytest.raises(TypeError):
            # got multiple values for keyword argument 'x'
            f(5.0, x=9)

    def test_state_access(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [x, In(a, value=1.0, name="a"), In(s, value=0.0, update=s + a * x)],
            s + a * x,
        )

        assert f[a] == 1.0
        assert f[s] == 0.0

        assert f(3.0) == 3.0
        assert f(3.0, a=2.0) == 9.0  # 3.0 + 2*3.0

        assert (
            f[a] == 1.0
        )  # state hasn't changed permanently, we just overrode it last line
        assert f[s] == 9.0

        f[a] = 5.0
        assert f[a] == 5.0
        assert f(3.0) == 24.0  # 9 + 3*5
        assert f[s] == 24.0

    def test_same_names(self):
        a, x, s = scalars("xxx")
        # implicit names would cause error.  What do we do?
        f = function([a, x, s], a + x + s)
        assert f(1, 2, 3) == 6
        with pytest.raises(TypeError):
            f(1, 2, x=3)

    def test_weird_names(self):
        a, x, s = scalars("xxx")

        with pytest.raises(TypeError):
            function([In(a, name=[])], [])

        def t():
            f = function(
                [
                    In(a, name={"adsf", ()}, value=1.0),
                    In(x, name=(), value=2.0),
                    In(s, name=scalar(), value=3.0),
                ],
                a + x + s,
            )
            return f

        with pytest.raises(TypeError):
            t()

    def test_copy(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )

        g = copy.copy(f)
        # if they both return, assume  that they return equivalent things.

        assert g.container[x].storage is not f.container[x].storage
        assert g.container[a].storage is not f.container[a].storage
        assert g.container[s].storage is not f.container[s].storage

        assert g.value[a] is f.value[a]  # should not have been copied
        assert (
            g.value[s] is not f.value[s]
        )  # should have been copied because it is mutable.
        assert not (g.value[s] != f.value[s]).any()  # its contents should be identical

        assert f(2, 1) == g(
            2
        )  # they should be in sync, default value should be copied.
        assert f(2, 1) == g(
            2
        )  # they should be in sync, default value should be copied.
        f(1, 2)  # put them out of sync
        assert f(1, 2) != g(1, 2)  # they should not be equal anymore.

    def test_copy_share_memory(self):
        x = fscalar("x")
        # SharedVariable for tests, one of them has update
        y = aesara.shared(value=1)
        z = aesara.shared(value=2)
        out = tanh((x + y + 2) / (x + z - 0.2) ** 2)

        # Test for different linkers
        for mode in ["FAST_RUN", "FAST_COMPILE"]:
            ori = aesara.function([x], [out], mode=mode, updates={z: z + 1})
            cpy = ori.copy(share_memory=True)

            # Test if memories shared
            storage_map_ori = ori.fn.storage_map
            storage_map_cpy = cpy.fn.storage_map
            fgraph_cpy = cpy.maker.fgraph

            # Assert intermediate and Constants storages are shared.
            # and output stoarges are not shared
            i_o_variables = fgraph_cpy.inputs + fgraph_cpy.outputs
            ori_storages = storage_map_ori.values()
            l = [
                val
                for key, val in storage_map_cpy.items()
                if key not in i_o_variables or isinstance(key, Constant)
            ]
            for storage in l:
                assert any([storage is s for s in ori_storages])

            # Assert storages of SharedVariable without updates are shared
            for (input, _1, _2), here, there in zip(
                ori.indices, ori.input_storage, cpy.input_storage
            ):
                assert here.data is there.data

    def test_swap_SharedVariable(self):
        i = iscalar()
        x_list = aesara.shared(value=np.random.rand(10).astype(config.floatX))

        x = scalar("x")
        # SharedVariable for tests, one of them has update
        y = aesara.shared(value=1, name="y")
        z = aesara.shared(value=2, name="z")
        m = aesara.shared(value=0, name="m")

        # SharedVariable to replace
        y_rpl = aesara.shared(value=3, name="y_rpl")
        z_rpl = aesara.shared(value=4, name="z_rpl")
        swap = {y: y_rpl, z: z_rpl}
        map_SV = {"y_rpl": y_rpl, "z_rpl": z_rpl}

        out = x + y + z + m

        # Test for different linkers
        # for mode in ["FAST_RUN","FAST_COMPILE"]:
        second_time = False
        for mode in ["FAST_RUN", "FAST_COMPILE"]:
            ori = aesara.function(
                [i],
                [out],
                mode=mode,
                updates=[(z, z + 1), (m, m + 2)],
                givens={x: x_list[i]},
            )
            cpy = ori.copy(swap=swap)

            # run fuction several time
            ori(1), cpy(1), cpy(2)

            # assert same SharedVariable are update in different function
            if not second_time:
                # m should be updated 3 times
                assert m.get_value() == 6
                # z should be updated once
                assert z.get_value() == 3
                # z_rpl should be updated twice
                assert z_rpl.get_value() == 6
                # y and y_rpl should not be updated
                assert y_rpl.get_value() == 3
                assert y.get_value() == 1
            elif second_time:
                # doule update for sharedvariable
                assert m.get_value() == 12
                assert z.get_value() == 4
                assert z_rpl.get_value() == 8
                assert y_rpl.get_value() == 3

            # test cpy function:
            # 2. SharedVariable is updatable -> values did update(z == 5)
            # 1. sharedvariable is swap ->  Rpl sharedvariables share storage
            names = map_SV.keys()
            for key in cpy.fn.storage_map:
                if key.name in names:
                    assert (
                        map_SV[key.name].container.storage[0]
                        == cpy.fn.storage_map[key][0]
                    )

            second_time = True

    def test_swap_SharedVariable_with_given(self):
        # A special testcase for logistic_sgd.py in Deep Learning Tutorial
        # This test assert that SharedVariable in different function have same storage

        train_x = aesara.shared(value=np.random.rand(10, 10).astype(config.floatX))
        test_x = aesara.shared(value=np.random.rand(10, 10).astype(config.floatX))

        train_y = aesara.shared(value=np.random.rand(10, 1).astype(config.floatX))
        test_y = aesara.shared(value=np.random.rand(10, 1).astype(config.floatX))

        i = iscalar("index")
        x = vector("x")
        y = vector("y")
        # this formular has no sense but for a test
        out = (aet_sum(x) - y) ** 2
        train = aesara.function(
            [i],
            out,
            givens={x: train_x[i], y: train_y[i]},
            updates={train_x: train_x + 0.1},
        )

        test_def = aesara.function([i], out, givens={x: test_x[i], y: test_y[i]})
        test_cpy = train.copy(
            swap={train_x: test_x, train_y: test_y}, delete_updates=True
        )

        for in1, in2 in zip(test_def.maker.inputs, test_cpy.maker.inputs):
            assert in1.value is in2.value

    def test_copy_delete_updates(self):
        w = iscalar("w")
        x = fscalar("x")
        # SharedVariable for tests, one of them has update
        y = aesara.shared(value=1, name="y")
        z = aesara.shared(value=2, name="z")
        out = x + y + z

        # Test for different linkers
        # for mode in ["FAST_RUN","FAST_COMPILE"]:
        # second_time = False
        for mode in ["FAST_RUN", "FAST_COMPILE"]:
            ori = aesara.function([x], out, mode=mode, updates={z: z * 2})
            cpy = ori.copy(delete_updates=True)

            assert cpy(1)[0] == 4
            assert cpy(1)[0] == 4
            assert cpy(1)[0] == 4

        # Test if unused implicit and explicit inputs from delete_updates
        # are ignored as intended.
        for mode in ["FAST_RUN", "FAST_COMPILE"]:
            ori = aesara.function([x], x, mode=mode, updates={z: z * 2})
            cpy = ori.copy(delete_updates=True)

            ori = aesara.function([x, w], x, mode=mode, updates={z: z + w})
            cpy = ori.copy(delete_updates=True)

    def test_shared_state0(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )
        g = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=f.container[s], update=s - a * x, mutable=True),
            ],
            s + a * x,
        )

        f(1, 2)
        assert f[s] == 2
        assert g[s] == 2
        g(1, 2)
        assert f[s] == 0
        assert g[s] == 0

    def test_shared_state1(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )
        g = function(
            [x, In(a, value=1.0, name="a"), In(s, value=f.container[s])], s + a * x
        )

        f(1, 2)
        assert f[s] == 2
        assert g[s] == 2
        f(1, 2)
        g(1, 2)
        assert f[s] == 4
        assert g[s] == 4

    def test_shared_state2(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=False),
            ],
            s + a * x,
        )
        g = function(
            [x, In(a, value=1.0, name="a"), In(s, value=f.container[s])], s + a * x
        )

        f(1, 2)
        assert f[s] == 2
        assert g[s] == 2
        f(1, 2)
        assert f[s] == 4
        assert g[s] == 4
        g(1, 2)  # has no effect on state
        assert f[s] == 4
        assert g[s] == 4

    def test_shared_state_not_implicit(self):
        # This test is taken from the documentation in
        # doc/topics/function.txt. If it does not pass anymore and yet the
        # behavior is still intended the doc and the test should both be
        # updated accordingly.
        x, s = scalars("xs")
        inc = function([x, In(s, update=(s + x), value=10.0)], [])
        dec = function(
            [x, In(s, update=(s - x), value=inc.container[s], implicit=False)], []
        )
        assert dec[s] is inc[s]
        inc[s] = 2
        assert dec[s] == 2
        dec(1)
        assert inc[s] == 1
        dec(1, 0)
        assert inc[s] == -1
        assert dec[s] == -1

    def test_constant_output(self):
        # Test that if the output is a constant, we respect the aesara memory interface
        f = aesara.function([], aet.constant([4]))
        # print f.maker.fgraph.toposort()
        out = f()
        assert (out == 4).all()
        out[0] = 3
        out2 = f()
        # If the following 2 asserts fail it mean Aesara broke it's memory contract.
        assert out2 is not out
        assert (out2 == 4).all()

        # Test that if the output is a constant and borrow, we respect the aesara memory interface
        f = aesara.function([], Out(aet.constant([4]), borrow=True))
        # print f.maker.fgraph.toposort()
        out = f()
        assert (out == 4).all()
        out[0] = 3
        out2 = f()

        if isinstance(
            aesara.compile.mode.get_default_mode(), aesara.compile.debugmode.DebugMode
        ):
            # In DebugMode, we don't implement optimization based on borrow on the output.
            assert (out2 == 4).all()
        else:
            assert out2 is out
            assert (out2 == 3).all()

    def test_borrow_input(self):
        # Tests that the contract for io.In is respected. When borrow=False, it should be
        # impossible for outputs to be aliased to the input variables provided by the user,
        # either through a view-map or a destroy map. New tests should be added in the future
        # when borrow=True is implemented.

        a = dmatrix()
        aval = np.random.rand(3, 3)

        # when borrow=False, test that a destroy map cannot alias output to input
        f = aesara.function([In(a, borrow=False)], Out(a + 1, borrow=True))
        assert np.all(f(aval) == aval + 1)
        assert not np.may_share_memory(aval, f(aval))

        # when borrow=False, test that a viewmap cannot alias output to input
        f = aesara.function([In(a, borrow=False)], Out(a[0, :], borrow=True))
        assert np.all(f(aval) == aval[0, :])
        assert not np.may_share_memory(aval, f(aval))

    def test_borrow_output(self):
        a = dmatrix()
        f = function([a], Out(a, borrow=False))
        o = np.ones((3, 3))
        assert o is not f(o)  # function no longer permits aliasing outputs to inputs

        f = function([a], Out(a * 4, borrow=False))
        o = np.ones((3, 3))
        four = f(o)
        assert np.all(four == 4)
        f(o + 0.1)  # should not clobber the memory used to store four
        assert np.all(four == 4)

        f = function([a], Out(a * 4, borrow=True), mode=Mode("c|py_nogc", "fast_run"))
        o = np.ones((3, 3))
        four = f(o)
        assert np.all(four == 4)
        f(o + 0.1)  # should clobber the memory used to store four
        if config.cxx:
            assert not np.all(four == 4)
        else:
            # The Elemwise.perform method don't reuse memory
            # as some numpy version don't support that correctly.
            assert np.all(four == 4)

    def test_disconnected_input(self):
        a = scalar("a")
        v = vector("v")
        with pytest.raises(UnusedInputError):
            function([a, v], v * 2)

        function([a, v], v * 2, on_unused_input="ignore")

    def test_masked_input(self):
        m = matrix("m")
        mt = m.T
        mt.name = "m.T"
        with pytest.raises(UnusedInputError):
            function([m, mt], mt * 2)
        function([m, mt], mt * 2, on_unused_input="ignore")

    def test_givens_input_var(self):
        # Ensure error is raised when trying to replace an input variable.

        x = scalar("x")
        y = x * 2
        with pytest.raises(RuntimeError):
            function([x], y, givens={x: x + 1})

    def test_free(self):
        # Make test on free() function

        x = vector("x")
        func = function([x], x + 1)
        func.fn.allow_gc = False
        func([1])

        check_list = []
        for key, val in func.fn.storage_map.items():
            if not isinstance(key, Constant):
                check_list.append(val)
        assert any([val[0] for val in check_list])

        func.free()

        for key, val in func.fn.storage_map.items():
            if not isinstance(key, Constant):
                assert val[0] is None

    def test_default_values(self):
        # Check that default values are restored
        # when an exception occurs in interactive mode.

        a, b = dscalars("a", "b")
        c = a + b
        func = aesara.function([In(a, name="first"), In(b, value=1, name="second")], c)
        x = func(first=1)
        try:
            func(second=2)
        except TypeError:
            assert func(first=1) == x

    def test_check_for_aliased_inputs(self):
        b = np.random.rand(5, 4)
        s1 = aesara.shared(b)
        s2 = aesara.shared(b)
        x1 = vector()

        # Assert cases we should not check for aliased inputs
        for d in [
            dict(outputs=[s1 + 1]),
            dict(outputs=[s1 + 1, s2 + 3]),
            dict(outputs=[s1 + 1], updates=[(s2, s2 + 3)]),
            dict(inputs=[x1], outputs=[x1 + 1], updates=[(s2, s2 + 3)]),
        ]:
            if "inputs" not in d:
                d["inputs"] = []
            f = aesara.function(**d)
            assert not f._check_for_aliased_inputs, d

        # Assert cases we should check for aliased inputs
        for d in [
            dict(
                inputs=[In(x1, borrow=True)],
                outputs=[x1 + 1],
                updates=[(s2, s2 + 3)],
            ),
            dict(
                inputs=[In(x1, borrow=True, mutable=True)],
                outputs=[x1 + 1],
                updates=[(s2, s2 + 3)],
            ),
            dict(
                inputs=[In(x1, mutable=True)],
                outputs=[x1 + 1],
                updates=[(s2, s2 + 3)],
            ),
        ]:
            if "inputs" not in d:
                d["inputs"] = []
            f = aesara.function(**d)

            assert f._check_for_aliased_inputs, d


class TestPicklefunction:
    def test_deepcopy(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )
        try:
            g = copy.deepcopy(f)
        except NotImplementedError as e:
            if e[0].startswith("DebugMode is not picklable"):
                return
            else:
                raise
        # if they both return, assume  that they return equivalent things.
        # print [(k,id(k)) for k in f.finder.keys()]
        # print [(k,id(k)) for k in g.finder.keys()]

        assert g.container[0].storage is not f.container[0].storage
        assert g.container[1].storage is not f.container[1].storage
        assert g.container[2].storage is not f.container[2].storage
        assert x not in g.container
        assert x not in g.value
        assert len(f.defaults) == len(g.defaults)
        assert f._check_for_aliased_inputs is g._check_for_aliased_inputs
        assert f.name == g.name
        assert f.maker.fgraph.name == g.maker.fgraph.name
        # print 'f.defaults = %s' % (f.defaults, )
        # print 'g.defaults = %s' % (g.defaults, )
        for ((f_req, f_feed, f_val), (g_req, g_feed, g_val)) in zip(
            f.defaults, g.defaults
        ):
            assert f_req == g_req and f_feed == g_feed and f_val == g_val

        assert g.value[1] is not f.value[1]  # should not have been copied
        assert (
            g.value[2] is not f.value[2]
        )  # should have been copied because it is mutable.
        assert not (g.value[2] != f.value[2]).any()  # its contents should be identical

        assert f(2, 1) == g(
            2
        )  # they should be in sync, default value should be copied.
        assert f(2, 1) == g(
            2
        )  # they should be in sync, default value should be copied.
        f(1, 2)  # put them out of sync
        assert f(1, 2) != g(1, 2)  # they should not be equal anymore.
        g(1, 2)  # put them back in sync
        assert f(3) == g(3)  # They should be in sync again.

    def test_deepcopy_trust_input(self):
        a = dscalar()  # the a is for 'anonymous' (un-named).
        x, s = dscalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )
        f.trust_input = True
        try:
            g = copy.deepcopy(f)
        except NotImplementedError as e:
            if e[0].startswith("DebugMode is not picklable"):
                return
            else:
                raise
        assert f.trust_input is g.trust_input
        f(np.asarray(2.0))
        with pytest.raises(
            (ValueError, AttributeError, aesara.compile.debugmode.InvalidValueError)
        ):
            f(2.0)
        g(np.asarray(2.0))
        with pytest.raises(
            (ValueError, AttributeError, aesara.compile.debugmode.InvalidValueError)
        ):
            g(2.0)

    def test_output_keys(self):
        x = vector()
        f = aesara.function([x], {"vec": x ** 2})
        o = f([2, 3, 4])
        assert isinstance(o, dict)
        assert np.allclose(o["vec"], [4, 9, 16])
        g = copy.deepcopy(f)
        o = g([2, 3, 4])
        assert isinstance(o, dict)
        assert np.allclose(o["vec"], [4, 9, 16])

    def test_deepcopy_shared_container(self):
        # Ensure that shared containers remain shared after a deep copy.
        a, x = scalars("ax")

        h = function([In(a, value=0.0)], a)
        f = function([x, In(a, value=h.container[a], implicit=True)], x + a)

        try:
            memo = {}
            ac = copy.deepcopy(a)
            memo.update({id(a): ac})
            hc = copy.deepcopy(h, memo=memo)
            memo.update({id(h): hc})
            fc = copy.deepcopy(f, memo=memo)
        except NotImplementedError as e:
            if e[0].startswith("DebugMode is not picklable"):
                return
            else:
                raise
        h[a] = 1
        hc[ac] = 2
        assert f[a] == 1
        assert fc[ac] == 2

    def test_pickle(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")

        f = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )

        try:
            # Note that here we also test protocol 0 on purpose, since it
            # should work (even though one should not use it).
            g = pickle.loads(pickle.dumps(f, protocol=0))
            g = pickle.loads(pickle.dumps(f, protocol=-1))
        except NotImplementedError as e:
            if e[0].startswith("DebugMode is not picklable"):
                return
            else:
                raise
        # if they both return, assume  that they return equivalent things.
        # print [(k,id(k)) for k in f.finder.keys()]
        # print [(k,id(k)) for k in g.finder.keys()]

        assert g.container[0].storage is not f.container[0].storage
        assert g.container[1].storage is not f.container[1].storage
        assert g.container[2].storage is not f.container[2].storage
        assert x not in g.container
        assert x not in g.value

        assert g.value[1] is not f.value[1]  # should not have been copied
        assert (
            g.value[2] is not f.value[2]
        )  # should have been copied because it is mutable.
        assert not (g.value[2] != f.value[2]).any()  # its contents should be identical

        assert f(2, 1) == g(
            2
        )  # they should be in sync, default value should be copied.
        assert f(2, 1) == g(
            2
        )  # they should be in sync, default value should be copied.
        f(1, 2)  # put them out of sync
        assert f(1, 2) != g(1, 2)  # they should not be equal anymore.

    def test_optimizations_preserved(self):
        a = dvector()  # the a is for 'anonymous' (un-named).
        x = dvector("x")
        s = dvector("s")
        xm = dmatrix("x")
        sm = dmatrix("s")

        f = function(
            [a, x, s, xm, sm],
            ((a.T.T) * (dot(xm, (sm.T.T.T)) + x).T * (x / x) + s),
        )
        old_default_mode = config.mode
        old_default_opt = config.optimizer
        old_default_link = config.linker
        try:
            try:
                str_f = pickle.dumps(f, protocol=-1)
                config.mode = "Mode"
                config.linker = "py"
                config.optimizer = "None"
                g = pickle.loads(str_f)
                # print g.maker.mode
                # print compile.mode.default_mode
            except NotImplementedError as e:
                if e[0].startswith("DebugMode is not pickl"):
                    g = "ok"
        finally:
            config.mode = old_default_mode
            config.optimizer = old_default_opt
            config.linker = old_default_link

        if g == "ok":
            return

        assert f.maker is not g.maker
        assert f.maker.fgraph is not g.maker.fgraph
        tf = f.maker.fgraph.toposort()
        tg = f.maker.fgraph.toposort()
        assert len(tf) == len(tg)
        for nf, ng in zip(tf, tg):
            assert nf.op == ng.op
            assert len(nf.inputs) == len(ng.inputs)
            assert len(nf.outputs) == len(ng.outputs)
            assert [i.type for i in nf.inputs] == [i.type for i in ng.inputs]
            assert [i.type for i in nf.outputs] == [i.type for i in ng.outputs]

    def test_multiple_functions(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")
        v = vector("v")

        # put in some inputs
        list_of_things = [s, x, v]

        # some derived thing, whose inputs aren't all in the list
        list_of_things.append(a * x + s)

        f1 = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )
        list_of_things.append(f1)

        # now put in a function sharing container with the previous one
        f2 = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=f1.container[s], update=s + a * x, mutable=True),
            ],
            s + a * x,
        )
        list_of_things.append(f2)

        assert isinstance(f2.container[s].storage, list)
        assert f2.container[s].storage is f1.container[s].storage

        # now put in a function with non-scalar
        v_value = np.asarray([2, 3, 4.0], dtype=config.floatX)
        f3 = function([x, In(v, value=v_value)], x + v)
        list_of_things.append(f3)

        # try to pickle the entire things
        try:
            saved_format = pickle.dumps(list_of_things, protocol=-1)
            new_list_of_things = pickle.loads(saved_format)
        except NotImplementedError as e:
            if e[0].startswith("DebugMode is not picklable"):
                return
            else:
                raise

        # now test our recovered new_list_of_things
        # it should be totally unrelated to the original
        # it should be interdependent in the same way as the original

        ol = list_of_things
        nl = new_list_of_things

        for i in range(4):
            assert nl[i] != ol[i]
            assert nl[i].type == ol[i].type
            assert nl[i].type is not ol[i].type

        # see if the implicit input got stored
        assert ol[3].owner.inputs[1] is s
        assert nl[3].owner.inputs[1] is not s
        assert nl[3].owner.inputs[1].type == s.type

        # moving on to the functions...
        for i in range(4, 7):
            assert nl[i] != ol[i]

        # looking at function number 1, input 's'
        assert nl[4][nl[0]] is not ol[4][ol[0]]
        assert nl[4][nl[0]] == ol[4][ol[0]]
        assert nl[4](3) == ol[4](3)

        # looking at function number 2, input 's'
        # make sure it's shared with the first function
        assert ol[4].container[ol[0]].storage is ol[5].container[ol[0]].storage
        assert nl[4].container[nl[0]].storage is nl[5].container[nl[0]].storage
        assert nl[5](3) == ol[5](3)
        assert nl[4].value[nl[0]] == 6

        assert np.all(nl[6][nl[2]] == np.asarray([2, 3.0, 4]))

    def test_broken_pickle_with_shared(self):
        saves = []

        def pers_save(obj):
            if isinstance(obj, np.ndarray):
                saves.append(obj)
                return len(saves) - 1
            else:
                return None

        def pers_load(id):
            return saves[id]

        b = np.random.rand(5, 4)

        x = matrix()
        y = aesara.shared(b)

        f = aesara.function([x], dot(x, y))

        from io import BytesIO

        fp = BytesIO()
        p = pickle.Pickler(fp, 2)
        p.persistent_id = pers_save
        try:
            p.dump(f)
        except NotImplementedError as e:
            if exc_message(e).startswith("DebugMode is not picklable"):
                return
            else:
                raise
        fp2 = BytesIO(fp.getvalue())
        fp.close()
        p = pickle.Unpickler(fp2)
        p.persistent_load = pers_load
        p.load()
        fp2.close()

    def test_pickle_class_with_functions(self):

        blah = SomethingToPickle()
        assert blah.f2.container[blah.s].storage is blah.f1.container[blah.s].storage

        try:
            blah2 = copy.deepcopy(blah)
        except NotImplementedError as e:
            if e[0].startswith("DebugMode is not picklable"):
                return
            else:
                raise

        assert (
            blah2.f2.container[blah2.s].storage is blah2.f1.container[blah2.s].storage
        )

        assert blah.f1[blah.s] == blah2.f1[blah2.s]

        blah.f2(5)
        assert blah.f1[blah.s] != blah2.f1[blah2.s]


class SomethingToPickle:
    def __init__(self):
        a = scalar()  # the a is for 'anonymous' (un-named).
        x, s = scalars("xs")
        v = vector("v")

        self.s = s
        self.x = x
        self.v = v

        self.e = a * x + s

        self.f1 = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=0.0, update=s + a * x, mutable=True),
            ],
            s + a * x,
        )

        self.f2 = function(
            [
                x,
                In(a, value=1.0, name="a"),
                In(s, value=self.f1.container[s], update=s + a * x, mutable=True),
            ],
            s + a * x,
        )


def test_empty_givens_updates():
    # Regression test for bug fixed in 8625e03.

    # Empty givens / updates dictionaries were not properly detected before,
    # triggering useless crashes at compile time.
    x = scalar()
    y = x * 2
    function([In(x)], y, givens={})
    function([In(x)], y, updates={})


@pytest.mark.skipif(
    not aesara.gpuarray.pygpu_activated or config.mode == "DEBUG_MODE",
    reason="DEBUG_MODE forces synchronous behaviour which breaks this test",
)
def test_sync_update():
    # This test if sync_update work. This can only be tested when
    # there is a GPU.  To test if we really sync, we compare a case we
    # can run in parallel GPU and CPU computation. Then we sync to
    # disable that parallel computation. Then we assert the time is
    # higher.

    # this import needs to go first because it generates the
    # local 'aesara' variable.  You get an UnboundLocalError otherwise.
    import tests.gpuarray.config

    sizes = [100, 500, 1000, 2000, 5000, 10000, 20000, 40000]
    size = sizes[0]
    w = aesara.gpuarray.gpuarray_shared_constructor(
        np.random.rand(size, size).astype("float32"),
        "w",
        target=tests.gpuarray.config.test_ctx_name,
    )
    x = aesara.gpuarray.gpuarray_shared_constructor(
        np.random.rand(size, size).astype("float32"),
        "x",
        target=tests.gpuarray.config.test_ctx_name,
    )

    updates = [(w, w + np.asarray(0.001, "float32") * dot(x, x))]

    f = aesara.function([], updates=updates, mode=tests.gpuarray.config.mode_with_gpu)
    assert len(f.maker.fgraph.apply_nodes) == 1
    assert any(
        isinstance(n.op, aesara.gpuarray.blas.GpuGemm)
        for n in f.maker.fgraph.apply_nodes
    )
    # Make sure libgpuarray have compile all kernels
    f()
    f.sync_shared()

    # Find a good size that will take about .5s.
    # This is to make the test more stable across different GPUs.
    size = sizes[-1]
    for i in sizes:
        data = np.random.rand(i, i).astype("float32")
        w.set_value(data)
        x.set_value(data)
        t0 = time.time()
        f()
        f.sync_shared()
        t1 = time.time()
        if (t1 - t0) < 0.5:
            continue
        size = i
        break
    # sync to make sure all computation are done
    f.sync_shared()

    t_0 = time.time()
    for i in range(3):
        f()
        # Sync after each call to see the slowdown from sync.
        f.sync_shared()
        time.sleep(0.5)
    t_1 = time.time()
    for i in range(3):
        f()
        time.sleep(0.5)
    f.sync_shared()
    # Sync to make sure all computation are finished.
    t_2 = time.time()
    d1 = t_1 - t_0
    d2 = t_2 - t_1
    assert d1 > d2, (d1, d2)


def test_FunctionMaker_cache_optimizations():

    opt_db_file = os.path.join(config.compiledir, "optimized_graphs.pkl")
    if os.path.exists(opt_db_file):
        os.remove(opt_db_file)

    floatX = "float32"
    mode = config.mode
    if mode in ["DEBUG_MODE", "DebugMode"]:
        mode = "FAST_RUN"

    graph_db_file = os.path.join(config.compiledir, "optimized_graphs.pkl")
    assert not os.path.exists(graph_db_file)

    with config.change_flags(cache_optimizations=True):
        a = fmatrix("a")
        b = fmatrix("b")
        c = aesara.shared(np.ones((10, 10), dtype=floatX))
        d = aesara.shared(np.ones((10, 10), dtype=floatX))
        e = aet_sum(aet_sum(aet_sum(a ** 2 + b) + c) + d)
        f1 = aesara.function([a, b], e, mode=mode)

        # FIXME: We can do much better about testing this.
        assert os.path.exists(graph_db_file)

        m = fmatrix("x1")
        n = fmatrix("x2")
        p = aesara.shared(np.ones((10, 10), dtype=floatX))
        q = aesara.shared(np.ones((10, 10), dtype=floatX))
        j = aet_sum(aet_sum(aet_sum(m ** 2 + n) + p) + q)
        f2 = aesara.function([m, n], j, mode=mode)

        in1 = np.ones((10, 10), dtype=floatX)
        in2 = np.ones((10, 10), dtype=floatX)
        assert f1(in1, in2) == f2(in1, in2)
