import sys

import numpy as np

import pytest

import theano
import theano.tensor

from six import reraise

from theano import config
from theano import gof
from theano.compile import debugmode

from tests import unittest_tools as utt


def test_debugmode_basic():
    x = theano.tensor.dvector()
    f = theano.function([x], ((2.0 * x) + 7) / 2.0, mode=debugmode.DebugMode())
    f([1, 2])


class BROKEN_ON_PURPOSE_Add(gof.Op):

    __props__ = ("py_offset",)

    def __init__(self, py_offset):
        gof.Op.__init__(self)
        self.py_offset = py_offset

    def make_node(self, a, b):
        a = theano.tensor.as_tensor_variable(a)
        b = theano.tensor.as_tensor_variable(b)
        assert a.type.dtype == "float64"
        assert a.type.dtype == b.type.dtype
        assert a.type.ndim == 1
        r = gof.Apply(self, [a, b], [a.type()])
        return r

    def perform(self, node, inp, out_):
        a, b = inp
        (out,) = out_
        z = a + b
        # ERROR TO ADD THIS CRAPPY OFFSET
        if self.py_offset:
            out[0] = z + 0.5
        else:
            out[0] = z

    def c_code_cache_version(self):
        return (1,)

    def c_code(self, node, name, inp, out, sub):
        a, b = inp
        (z,) = out
        return """
        if (PyArray_NDIM(%(a)s) != 1) {PyErr_SetString(PyExc_NotImplementedError, "rank(a) != 1"); %(fail)s;}
        if (PyArray_NDIM(%(b)s) != 1) {PyErr_SetString(PyExc_NotImplementedError, "rank(b) != 1"); %(fail)s;}

        if (PyArray_DESCR(%(a)s)->type_num != NPY_DOUBLE)
        {PyErr_SetString(PyExc_NotImplementedError, "a dtype not NPY_DOUBLE"); %(fail)s;}

        if (PyArray_DESCR(%(b)s)->type_num != NPY_DOUBLE)
        {PyErr_SetString(PyExc_NotImplementedError, "b's dtype not NPY_DOUBLE"); %(fail)s;}

        if (PyArray_DIMS(%(a)s)[0] != PyArray_DIMS(%(b)s)[0])
        {PyErr_SetString(PyExc_NotImplementedError, "a and b have different lengths"); %(fail)s;}

        if ((!%(z)s)
            || (PyArray_DIMS(%(z)s)[0] != PyArray_DIMS(%(b)s)[0])
            )
        {
            {Py_XDECREF(%(z)s);}
            npy_intp dims[] = {0};
            dims[0] = PyArray_DIMS(%(b)s)[0];
            %(z)s = (PyArrayObject*) PyArray_SimpleNew(1, dims, PyArray_DESCR(%(b)s)->type_num);
        }

        {
            for (npy_intp m = 0; m < PyArray_DIMS(%(z)s)[0]; ++m)
            {
                ((double*)PyArray_GETPTR1(%(z)s, m))[0]
                = 0.5
                + ((double*)PyArray_GETPTR1(%(a)s, m))[0]
                + ((double*)PyArray_GETPTR1(%(b)s, m))[0] ;
            }
        }
        """ % dict(
            locals(), **sub
        )


# inconsistent is a invalid op, whose perform and c_code do not match
inconsistent = BROKEN_ON_PURPOSE_Add(False)

# off_by_half is a good op, that is different from theano.sparse.sd_csc
off_by_half = BROKEN_ON_PURPOSE_Add(True)


class WeirdBrokenOp(gof.Op):
    """
    This op can be inplace if behaviour is 'times1_inplace'
    This op can be destructive if behaviour is 'times2_inplace'

    In both cases, it does not set the destroy_map or view_map correctly so
    it should raise an error in DebugMode.
    """

    __props__ = ("behaviour",)

    def __init__(self, behaviour):
        gof.Op.__init__(self)
        self.behaviour = behaviour

    def make_node(self, a):
        a_ = theano.tensor.as_tensor_variable(a)
        r = gof.Apply(self, [a_], [a_.type()])
        return r

    def dontuse_perform(self, node, inp, out_):
        (a,) = inp
        (out,) = out_
        if self.behaviour == "times2":
            out[0] = a * 2
        elif self.behaviour == "times2_inplace":
            out[0] = a
            out[0] *= 2
        elif self.behaviour == "times1":
            out[0] = a * 1
        elif self.behaviour == "times1_inplace":
            out[0] = a
        else:
            raise ValueError(self.behaviour)

    def c_code_cache_version(self):
        return (1,)

    def c_code(self, node, name, inp, out, sub):
        (a,) = inp
        (z,) = out
        if "inplace" in self.behaviour:
            z_code = """
            {Py_XDECREF(%(z)s);}
            Py_INCREF(%(a)s);
            %(z)s = %(a)s;
            """
        else:
            z_code = """
            {Py_XDECREF(%(z)s);}
            %(z)s = (PyArrayObject*) PyArray_SimpleNew(1, PyArray_DIMS(%(a)s), PyArray_DESCR(%(a)s)->type_num);
            """
        prep_vars = """
            //the output array has size M x N
            npy_intp M = PyArray_DIMS(%(a)s)[0];
            npy_intp Sa = PyArray_STRIDES(%(a)s)[0] / PyArray_DESCR(%(a)s)->elsize;
            npy_intp Sz = PyArray_STRIDES(%(z)s)[0] / PyArray_DESCR(%(z)s)->elsize;

            npy_double * Da = (npy_double*)PyArray_BYTES(%(a)s);
            npy_double * Dz = (npy_double*)PyArray_BYTES(%(z)s);

            //clear the output array
            for (npy_intp m = 0; m < M; ++m)
            {
        """

        if self.behaviour == "times2":
            behaviour = "     Dz[m * Sz] = 2 * Da[m * Sa]; "
            # out[0] = a * 2
        elif self.behaviour == "times2_inplace":
            # out[0] = a
            # out[0] *= 2
            behaviour = "     Dz[m * Sz] = 2 * Da[m * Sa]; "
        elif self.behaviour == "times1":
            # out[0] = a * 1
            behaviour = "     Dz[m * Sz] = Da[m * Sa]; "
        elif self.behaviour == "times1_inplace":
            # out[0] = a
            behaviour = ""
        else:
            raise ValueError(self.behaviour)

        prep_vars2 = """
            }
        """

        total = (z_code + prep_vars + behaviour + prep_vars2) % dict(locals(), **sub)
        return total


wb2i = WeirdBrokenOp("times2_inplace")
wb2 = WeirdBrokenOp("times2")
wb1i = WeirdBrokenOp("times1_inplace")
wb1 = WeirdBrokenOp("times1")


@pytest.mark.skipif(
    not theano.config.cxx, reason="G++ not available, so we need to skip this test."
)
def test_badthunkoutput():
    # Check if the c and python code is consistent.
    a = theano.tensor.dvector()
    b = theano.tensor.dvector()

    f_good = theano.function(
        [a, b],
        off_by_half(a, b),
        mode=debugmode.DebugMode(check_c_code=theano.config.cxx),
    )
    f_inconsistent = theano.function(
        [a, b],
        inconsistent(a, b),
        mode=debugmode.DebugMode(check_c_code=theano.config.cxx),
    )

    # this should evaluate with no error
    f_good([1.0, 2.0, 3.0], [2, 3, 4])

    with pytest.raises(debugmode.BadThunkOutput) as einfo:
        f_inconsistent([1.0, 2.0, 3.0], [2, 3, 4])

    assert einfo.value.r.owner.op is inconsistent


def test_badoptimization():
    @gof.local_optimizer([theano.tensor.add])
    def insert_broken_add(node):
        if node.op == theano.tensor.add:
            return [off_by_half(*node.inputs)]
        return False

    edb = gof.EquilibriumDB()
    edb.register("insert_broken_add", insert_broken_add, "all")
    opt = edb.query("+all")

    a = theano.tensor.dvector()
    b = theano.tensor.dvector()

    f = theano.function([a, b], a + b, mode=debugmode.DebugMode(optimizer=opt))

    with pytest.raises(debugmode.BadOptimization) as einfo:
        f(
            [1.0, 2.0, 3.0],
            [2, 3, 4],
        )
    assert str(einfo.value.reason) == "insert_broken_add"


def test_badoptimization_opt_err():
    # This variant of test_badoptimization() replace the working code
    # with a new apply node that will raise an error.
    @gof.local_optimizer([theano.tensor.add])
    def insert_bigger_b_add(node):
        if node.op == theano.tensor.add:
            inputs = list(node.inputs)
            if inputs[-1].owner is None:
                inputs[-1] = theano.tensor.concatenate((inputs[-1], inputs[-1]))
                return [node.op(*inputs)]
        return False

    @gof.local_optimizer([theano.tensor.add])
    def insert_bad_dtype(node):
        if node.op == theano.tensor.add:
            inputs = list(node.inputs)
            if inputs[-1].owner is None:

                return [node.outputs[0].astype("float32")]
        return False

    edb = gof.EquilibriumDB()
    edb.register("insert_bigger_b_add", insert_bigger_b_add, "all")
    opt = edb.query("+all")
    edb2 = gof.EquilibriumDB()
    edb2.register("insert_bad_dtype", insert_bad_dtype, "all")
    opt2 = edb2.query("+all")

    a = theano.tensor.dvector()
    b = theano.tensor.dvector()

    f = theano.function([a, b], a + b, mode=debugmode.DebugMode(optimizer=opt))
    with pytest.raises(ValueError, match=r"insert_bigger_b_add"):
        f(
            [1.0, 2.0, 3.0],
            [2, 3, 4],
        )

    # Test that opt that do an illegal change still get the error from gof.
    with pytest.raises(
        theano.gof.toolbox.BadOptimization, match=r"insert_bad_dtype"
    ) as einfo:
        with theano.change_flags(on_opt_error="raise"):
            f2 = theano.function(
                [a, b],
                a + b,
                mode=debugmode.DebugMode(optimizer=opt2, stability_patience=1),
            )
        f2(
            [1.0, 2.0, 3.0],
            [2, 3, 4],
        )

    # Test that we can reraise the error with an extended message
    with pytest.raises(theano.gof.toolbox.BadOptimization):
        e = einfo.value
        new_e = e.__class__("TTT" + str(e))
        exc_type, exc_value, exc_trace = sys.exc_info()
        exc_value = new_e
        reraise(e.__class__, exc_value, exc_trace)


def test_stochasticoptimization():

    # this optimization alternates between triggering and not triggering.

    last_time_replaced = [False]

    @gof.local_optimizer([theano.tensor.add])
    def insert_broken_add_sometimes(node):
        if node.op == theano.tensor.add:
            last_time_replaced[0] = not last_time_replaced[0]
            if last_time_replaced[0]:
                return [off_by_half(*node.inputs)]
        return False

    edb = gof.EquilibriumDB()
    edb.register("insert_broken_add_sometimes", insert_broken_add_sometimes, "all")
    opt = edb.query("+all")

    a = theano.tensor.dvector()
    b = theano.tensor.dvector()

    with pytest.raises(debugmode.StochasticOrder):
        theano.function(
            [a, b],
            theano.tensor.add(a, b),
            mode=debugmode.DebugMode(
                optimizer=opt,
                check_c_code=True,
                stability_patience=max(2, config.DebugMode.patience),
            ),
        )


@pytest.mark.skipif(
    not theano.config.cxx, reason="G++ not available, so we need to skip this test."
)
def test_just_c_code():
    x = theano.tensor.dvector()
    f = theano.function([x], wb2(x), mode=debugmode.DebugMode(check_py_code=False))
    assert np.all(f([1, 2]) == [2, 4])


def test_baddestroymap():
    class BadAdd(gof.Op):
        def make_node(self, a, b):
            c = a.type()
            return gof.Apply(self, [a, b], [c])

        def perform(self, node, inp, out):
            a, b = inp
            (c,) = out
            c[0] = a
            c[0] += b

    x = theano.tensor.dvector()
    y = theano.tensor.dvector()
    f = theano.function([x, y], BadAdd()(x, y), mode="DEBUG_MODE")

    with pytest.raises(debugmode.BadDestroyMap):
        f([1, 2], [3, 4])


@pytest.mark.skipif(
    not theano.config.cxx, reason="G++ not available, so we need to skip this test."
)
def test_baddestroymap_c():
    x = theano.tensor.dvector()
    f = theano.function([x], wb2i(x), mode=debugmode.DebugMode(check_py_code=False))
    with pytest.raises(debugmode.BadDestroyMap):
        assert np.all(f([1, 2]) == [2, 4])


class TestViewMap:
    class BadAddRef(gof.Op):
        def make_node(self, a, b):
            c = b.type()
            return gof.Apply(self, [a, b], [c])

        def perform(self, node, inp, out):
            a, b = inp
            (c,) = out
            c[0] = b

    class BadAddSlice(gof.Op):
        def make_node(self, a, b):
            c = b.type()
            return gof.Apply(self, [a, b], [c])

        def perform(self, node, inp, out):
            a, b = inp
            (c,) = out
            c[0] = b[1:3]

    def test_badviewmap_ref(self):
        x = theano.tensor.dvector()
        y = theano.tensor.dvector()
        f = theano.function([x, y], self.BadAddRef()(x, y), mode="DEBUG_MODE")
        with pytest.raises(debugmode.BadViewMap):
            f([1, 2], [3, 4])

    def test_badviewmap_slice(self):
        x = theano.tensor.dvector()
        y = theano.tensor.dvector()
        f = theano.function([x, y], self.BadAddSlice()(x, y), mode="DEBUG_MODE")
        with pytest.raises(debugmode.BadViewMap):
            f([1, 2], [3, 4])

    def test_goodviewmap(self):
        goodop = self.BadAddRef()
        goodop.view_map = {0: [1]}
        x = theano.tensor.dvector()
        y = theano.tensor.dvector()
        f = theano.function([x, y], goodop(x, y), mode="DEBUG_MODE")
        # Shouldn't raise an error
        f([1, 5, 1], [3, 4, 2, 1, 4])

    @pytest.mark.skipif(
        not theano.config.cxx, reason="G++ not available, so we need to skip this test."
    )
    def test_badviewmap_c(self):
        x = theano.tensor.dvector()
        f = theano.function([x], wb1i(x), mode=debugmode.DebugMode(check_py_code=False))
        with pytest.raises(debugmode.BadViewMap):
            f([1, 2])

    def test_aliased_outputs_ok(self):
        # here aliased outputs is ok because they are both aliased to an input
        # as well
        class CustomOp(gof.Op):
            view_map = {0: [0], 1: [0]}

            def make_node(self, a, b):
                c = a.type()
                d = a.type()
                return gof.Apply(self, [a, b], [c, d])

            def perform(self, node, inp, out):
                a, b = inp
                c, d = out
                c[0] = a
                d[0] = a[1:]

        x = theano.tensor.dvector("x")
        y = theano.tensor.dvector("y")
        f = theano.function([x, y], CustomOp()(x, y), mode="DEBUG_MODE")

        r0, r1 = f([1, 2, 3, 4], [5, 6, 7, 8])

        assert np.all(r0 == [1, 2, 3, 4])
        assert np.all(r1 == [2, 3, 4])

    def test_aliased_outputs_ok_output(self):
        # here aliased outputs is ok because they are both outputs of the
        # function as a whole and thus not destroy-able
        class CustomOp(gof.Op):
            def make_node(self, a, b):
                c = a.type()
                d = a.type()
                return gof.Apply(self, [a, b], [c, d])

            def perform(self, node, inp, out):
                a, b = inp
                c, d = out
                r = a * 2
                c[0] = r
                d[0] = r[1:]

        x = theano.tensor.dvector()
        y = theano.tensor.dvector()
        f = theano.function([x, y], CustomOp()(x, y), mode="DEBUG_MODE")

        r0, r1 = f([1, 2, 3, 4], [5, 6, 7, 8])

        assert np.all(r0 == [2, 4, 6, 8])
        assert np.all(r1 == [4, 6, 8])

    def test_aliased_outputs_ok_shadow(self):
        # here the alias between outputs is ok because one of them is not used
        # for subsequent computation.  This is like the case where we use one
        # output as a memory buffer to serve another output.
        class CustomOp(gof.Op):
            def make_node(self, a, b):
                c = a.type()
                d = a.type()
                return gof.Apply(self, [a, b], [c, d])

            def perform(self, node, inp, out):
                a, b = inp
                c, d = out
                r = a * 1
                c[0] = r
                d[0] = r[1:]

        x = theano.tensor.dvector("x")
        y = theano.tensor.dvector("y")
        f = theano.function([x, y], CustomOp()(x, y)[0] * 2, mode="DEBUG_MODE")

        r0 = f([1, 2, 3, 4], [5, 6, 7, 8])

        assert np.all(r0 == [2, 4, 6, 8])

    def test_aliased_outputs_bad(self):
        # here the alias between outputs is not ok because destroying one
        # destroys the other, but there's no way to warn theano about it
        # through the view_map mechanism.
        class CustomOp(gof.Op):
            def make_node(self, a, b):
                c = a.type()
                d = a.type()
                return gof.Apply(self, [a, b], [c, d])

            def perform(self, node, inp, out):
                a, b = inp
                c, d = out
                r = a * 1
                c[0] = r[:-1]
                d[0] = r[1:]

        custom_op = CustomOp()

        x = theano.tensor.dvector()
        y = theano.tensor.dvector()
        bad_xy0, bad_xy1 = custom_op(x, y)
        out = bad_xy0 * 2 + bad_xy1 * 2
        f = theano.function([x, y], out, mode="DEBUG_MODE")

        with pytest.raises(debugmode.BadViewMap):
            f([1, 2, 3, 4], [5, 6, 7, 8])

        # the situation can be rescued by picking one of the inputs and
        # pretending that it is aliased to both the outputs.
        # This unfairly disables any destructive operations on the
        # input, but guarantees correctness.
        # custom_op.view_map = {0:[0], 1:[1]}
        # f([1,2,3,4],[5,6,7,8])


class TestCheckIsfinite:
    def setup_method(self):
        self.old_ts = theano.tensor.TensorType.filter_checks_isfinite
        self.old_dm = theano.compile.mode.predefined_modes["DEBUG_MODE"].check_isfinite

    def teardown_method(self):
        theano.tensor.TensorType.filter_checks_isfinite = self.old_ts
        theano.compile.mode.predefined_modes["DEBUG_MODE"].check_isfinite = self.old_dm

    def test_check_isfinite(self):
        x = theano.tensor.vector()
        f = theano.function([x], (x + 2) * 5, mode="DEBUG_MODE")
        g = theano.function([x], theano.tensor.log(x), mode="DEBUG_MODE")

        # this should work
        f(np.log([3, 4, 5]).astype(config.floatX))

        # if TensorType.filter_checks_isfinite were true, these would raise
        # ValueError
        # if not, DebugMode will check internally, and raise InvalidValueError
        # passing an invalid value as an input should trigger ValueError
        with pytest.raises(debugmode.InvalidValueError):
            f(np.log([3, -4, 5]).astype(config.floatX))
        with pytest.raises(debugmode.InvalidValueError):
            f((np.asarray([0, 1.0, 0]) / 0).astype(config.floatX))
        with pytest.raises(debugmode.InvalidValueError):
            f((np.asarray([1.0, 1.0, 1.0]) / 0).astype(config.floatX))

        # generating an invalid value internally should trigger
        # InvalidValueError
        with pytest.raises(debugmode.InvalidValueError):
            g(np.asarray([3, -4, 5], dtype=config.floatX))

        # this should disable the exception
        theano.tensor.TensorType.filter_checks_isfinite = False
        theano.compile.mode.predefined_modes["DEBUG_MODE"].check_isfinite = False
        # insert several Inf
        f(np.asarray(np.asarray([1.0, 1.0, 1.0]) / 0, dtype=config.floatX))

    def test_check_isfinite_disabled(self):
        x = theano.tensor.dvector()
        f = theano.function(
            [x], (x + 2) * 5, mode=debugmode.DebugMode(check_isfinite=False)
        )

        # nan should go through
        f(np.log([3, -4, 5]))

        # inf should go through
        infs = np.asarray([1.0, 1.0, 1.0]) / 0
        # print infs
        f(infs)
        return


class BrokenCImplementationAdd(gof.Op):

    __props__ = ()

    def make_node(self, a, b):
        a = theano.tensor.as_tensor_variable(a)
        b = theano.tensor.as_tensor_variable(b)
        assert a.type.dtype == "float32"
        assert a.type.dtype == b.type.dtype
        assert a.type.ndim == 2
        r = gof.Apply(self, [a, b], [a.type()])
        return r

    def perform(self, node, inp, out_):
        # print 'executing python perform'
        a, b = inp
        (out,) = out_
        z = a + b
        # print 'out[0] was:', out[0]
        out[0] = z

    def c_code_cache_version(self):
        return (1,)

    def c_code(self, node, name, inp, out, sub):
        a, b = inp
        (z,) = out
        debug = 0
        return """
        //printf("executing c_code\\n");
        if (PyArray_NDIM(%(a)s) != 2) {PyErr_SetString(PyExc_NotImplementedError, "rank(a) != 2"); %(fail)s;}
        if (PyArray_NDIM(%(b)s) != 2) {PyErr_SetString(PyExc_NotImplementedError, "rank(b) != 2"); %(fail)s;}

        if (PyArray_DESCR(%(a)s)->type_num != NPY_FLOAT)
        {PyErr_SetString(PyExc_NotImplementedError, "a dtype not NPY_FLOAT"); %(fail)s;}

        if (PyArray_DESCR(%(b)s)->type_num != NPY_FLOAT)
        {PyErr_SetString(PyExc_NotImplementedError, "b's dtype not NPY_FLOAT"); %(fail)s;}

        if (PyArray_DIMS(%(a)s)[0] != PyArray_DIMS(%(a)s)[1])
        {PyErr_SetString(PyExc_NotImplementedError, "a is not square"); %(fail)s;}

        if (PyArray_DIMS(%(b)s)[0] != PyArray_DIMS(%(b)s)[1])
        {PyErr_SetString(PyExc_NotImplementedError, "b is not square"); %(fail)s;}

        if (PyArray_DIMS(%(a)s)[0] != PyArray_DIMS(%(b)s)[0])
        {PyErr_SetString(PyExc_NotImplementedError, "a and b have different dimensions"); %(fail)s;}

        // We do not check for c_contiguous property here
        if (%(debug)s)
        {
            if (!%(z)s)
                printf("%(z)s is not there, %%p \\n", %(z)s);
            else if (PyArray_DIMS(%(z)s)[0] != PyArray_DIMS(%(b)s)[0])
                printf("Dimension 0 mismatch for %(z)s and %(b)s\\n");
            else if (PyArray_DIMS(%(z)s)[1] != PyArray_DIMS(%(b)s)[1])
                printf("Dimension 1 mismatch for %(z)s and %(b)s\\n");
            else
                printf("Reusing %(z)s\\n");
        }

        if ((!%(z)s)
            || (PyArray_DIMS(%(z)s)[0] != PyArray_DIMS(%(b)s)[0])
            || (PyArray_DIMS(%(z)s)[1] != PyArray_DIMS(%(b)s)[1])
            )
        {
            Py_XDECREF(%(z)s);
            npy_intp dims[] = {0, 0};
            dims[0] = PyArray_DIMS(%(b)s)[0];
            dims[1] = PyArray_DIMS(%(b)s)[1];
            %(z)s = (PyArrayObject*) PyArray_SimpleNew(2, dims, PyArray_DESCR(%(b)s)->type_num);
        }

        // Let us assume that %(z)s is c_contiguous
        {
            dtype_%(z)s * z = ((dtype_%(z)s*)(PyArray_GETPTR2(%(z)s,0,0)));
            for (int i=0; i<PyArray_DIMS(%(b)s)[0]; i++)
            {
                for (int j=0; j<PyArray_DIMS(%(b)s)[1]; j++)
                {
                    *z = ((float*)PyArray_GETPTR2(%(a)s, i, j))[0] +
                         ((float*)PyArray_GETPTR2(%(b)s, i, j))[0] ;
                    z++;
                }
            }
        }
        """ % dict(
            locals(), **sub
        )


class VecAsRowAndCol(gof.Op):
    """
    Transforms a vector into a row and a column.

    This Op exists to check everything is correct when an Op has
    two outputs with different broadcasting patterns.
    """

    __props__ = ()

    def make_node(self, v):
        if not isinstance(v, gof.Variable):
            v = theano.tensor.as_tensor_variable(v)
        assert v.type.ndim == 1
        type_class = type(v.type)
        out_r_type = type_class(dtype=v.dtype, broadcastable=(True, False))
        out_c_type = type_class(dtype=v.dtype, broadcastable=(False, True))
        return gof.Apply(self, [v], [out_r_type(), out_c_type()])

    def perform(self, node, inp, out):
        (v,) = inp
        r, c = out
        lv = v.shape[0]
        if (r[0] is None) or (r[0].shape != (1, lv)):
            r[0] = node.outputs[0].type.value_zeros((1, lv))

        if (c[0] is None) or (c[0].shape != (lv, 1)):
            c[0] = node.outputs[1].type.value_zeros((lv, 1))

        for i in range(lv):
            r[0][0, i] = v[i]
            c[0][i, 0] = v[i]


class TestPreallocatedOutput:
    def setup_method(self):
        self.rng = np.random.RandomState(seed=utt.fetch_seed())

    def test_f_contiguous(self):
        a = theano.tensor.fmatrix("a")
        b = theano.tensor.fmatrix("b")
        z = BrokenCImplementationAdd()(a, b)
        # In this test, we do not want z to be an output of the graph.
        out = theano.tensor.dot(z, np.eye(7))

        a_val = self.rng.randn(7, 7).astype("float32")
        b_val = self.rng.randn(7, 7).astype("float32")

        # Should work
        mode = debugmode.DebugMode(check_preallocated_output=["c_contiguous"])

        f = theano.function([a, b], out, mode=mode)
        f(a_val, b_val)
        # print 'out_val =', out_val
        # print out_val.strides

        # Should raise an Exception, since the output buffer is
        # used incorrectly.
        mode = debugmode.DebugMode(check_preallocated_output=["f_contiguous"])

        f = theano.function([a, b], out, mode=mode)

        if theano.config.cxx:
            with pytest.raises(debugmode.BadThunkOutput):
                f(a_val, b_val)
        else:
            # The python code of this op is good.
            f(a_val, b_val)

    def test_f_contiguous_out(self):
        # Same test as test_f_contiguous, but check that it works
        # even if z _is_ the output of the graph
        a = theano.tensor.fmatrix("a")
        b = theano.tensor.fmatrix("b")
        out = BrokenCImplementationAdd()(a, b)

        a_val = self.rng.randn(7, 7).astype("float32")
        b_val = self.rng.randn(7, 7).astype("float32")

        # Should work
        mode = debugmode.DebugMode(check_preallocated_output=["c_contiguous"])

        f = theano.function([a, b], out, mode=mode)
        f(a_val, b_val)
        # print 'out_val =', out_val
        # print out_val.strides

        # Should raise an Exception, since the output buffer is
        # used incorrectly.
        mode = debugmode.DebugMode(check_preallocated_output=["f_contiguous"])

        f = theano.function([a, b], out, mode=mode)

        if theano.config.cxx:
            with pytest.raises(debugmode.BadThunkOutput):
                f(a_val, b_val)
        else:
            # The python code of this op is good.
            f(a_val, b_val)

    def test_output_broadcast_tensor(self):
        v = theano.tensor.fvector("v")
        c, r = VecAsRowAndCol()(v)
        f = theano.function([v], [c, r])

        v_val = self.rng.randn(5).astype("float32")
        f(v_val)
