import unittest

from theano.gof import graph
from theano.gof.graph import Variable, Apply, Constant
from theano.gof.type import Type
from theano.gof.op import Op
from theano.gof import fg

from theano.gof.link import *
from theano.compat import cmp

def as_variable(x):
    assert isinstance(x, Variable)
    return x


class TDouble(Type):
    def filter(self, data):
        return float(data)

tdouble = TDouble()


def double(name):
    return Variable(tdouble, None, None, name=name)


class MyOp(Op):
    def __init__(self, nin, name, impl=None):
        self.nin = nin
        self.name = name
        if impl:
            self.impl = impl

    def make_node(self, *inputs):
        assert len(inputs) == self.nin
        inputs = map(as_variable, inputs)
        for input in inputs:
            if input.type is not tdouble:
                raise Exception("Error 1")
        outputs = [double(self.name + "_R")]
        return Apply(self, inputs, outputs)

    def __str__(self):
        return self.name

    def perform(self, node, inputs, out_):
        out, = out_
        out[0] = self.impl(*inputs)

add = MyOp(2, 'Add', lambda x, y: x + y)
sub = MyOp(2, 'Sub', lambda x, y: x - y)
mul = MyOp(2, 'Mul', lambda x, y: x * y)
div = MyOp(2, 'Div', lambda x, y: x / y)


def notimpl(self, x):
    raise NotImplementedError()


raise_err = MyOp(1, 'RaiseErr', notimpl)


def inputs():
    x = double('x')
    y = double('y')
    z = double('z')
    return x, y, z


def perform_linker(fgraph):
    lnk = PerformLinker().accept(fgraph)
    return lnk


def FunctionGraph(inputs, outputs):
    e = fg.FunctionGraph(inputs, outputs)
    return e


class TestPerformLinker(unittest.TestCase):
    def test_thunk(self):
        x, y, z = inputs()
        e = mul(add(x, y), div(x, y))
        fn, i, o = perform_linker(FunctionGraph([x, y, z], [e])).make_thunk()
        i[0].data = 1
        i[1].data = 2
        fn()
        assert o[0].data == 1.5

    def test_function(self):
        x, y, z = inputs()
        e = mul(add(x, y), div(x, y))
        fn = perform_linker(FunctionGraph([x, y, z], [e])).make_function()
        assert fn(1.0, 2.0, 3.0) == 1.5

    def test_constant(self):
        x, y, z = inputs()
        y = Constant(tdouble, 2.0)
        e = mul(add(x, y), div(x, y))
        fn = perform_linker(FunctionGraph([x], [e])).make_function()
        assert fn(1.0) == 1.5

    def test_input_output_same(self):
        x, y, z = inputs()
        fn = perform_linker(FunctionGraph([x], [x])).make_function()
        assert 1.0 is fn(1.0)

    def test_input_dependency0(self):
        x, y, z = inputs()
        a, d = add(x, y), div(x, y)
        e = mul(a, d)
        fn = perform_linker(FunctionGraph(*graph.clone([x, y, a], [e]))).make_function()
        assert fn(1.0, 2.0, 9.0) == 4.5

    def test_skiphole(self):
        x, y, z = inputs()
        a = add(x, y)
        r = raise_err(a)
        e = add(r, a)
        fn = perform_linker(FunctionGraph(*graph.clone([x, y, r], [e]))).make_function()
        assert fn(1.0, 2.0, 4.5) == 7.5


def wrap_linker(fgraph, linkers, wrapper):
    lnk = WrapLinker(linkers, wrapper).accept(fgraph)
    return lnk


class TestWrapLinker(unittest.TestCase):
    def test_0(self):
        nodes = []

        def wrap(i, node, th):
            nodes.append(node.op)

        x, y, z = inputs()
        e = mul(add(x, y), div(x, y))
        fn, i, o = wrap_linker(
                FunctionGraph([x, y, z], [e]),
                [PerformLinker(allow_gc=False)], wrap).make_thunk()
        i[0].data = 1
        i[1].data = 2
        fn()
        assert nodes == [div, add, mul]
        assert o[0].data is None

    def test_1(self):
        nodes = []

        def wrap(i, node, th):
            nodes.append(node.op)
            th()

        x, y, z = inputs()
        e = mul(add(x, y), div(x, y))
        fn, i, o = wrap_linker(
                FunctionGraph([x, y, z], [e]),
                [PerformLinker(allow_gc=False)], wrap).make_thunk()
        i[0].data = 1
        i[1].data = 2
        fn()
        assert nodes == [div, add, mul]
        assert o[0].data == 1.5

def test_sort_schedule_fn():
    import theano
    from theano.gof.sched import sort_schedule_fn, make_depends
    x = theano.tensor.matrix('x')
    y = theano.tensor.dot(x[:5]*2, x.T+1).T
    str_cmp = lambda a, b: cmp(str(a), str(b)) # lexicographical sort
    linker = theano.OpWiseCLinker(schedule=sort_schedule_fn(str_cmp))
    mode = theano.Mode(linker=linker)
    f = theano.function((x,), (y,), mode=mode)

    nodes = f.maker.linker.make_all()[-1]
    depends = make_depends()
    for a, b in zip(nodes[:-1], nodes[1:]):
        if not depends((b,a)):
            assert str(a) < str(b)
