import unittest

import modelflow.elementarynodes
import modelflow.graph
import modelflow.input
from modelflow import elementarynodes as el


class StaticTestComputation(modelflow.graph.StaticGraph):
        one = modelflow.input.Input()
        two = modelflow.input.Input()

        def construct_graph(self):
            plus = modelflow.elementarynodes.Plus()
            self.connect_input_to_node("one", plus, "arg1")
            self.connect_input_to_node("two", plus, "arg2")
            self.set_output(plus)


class StaticGraphofStaticGraph(modelflow.graph.StaticGraph):
        two = modelflow.input.Input()

        def construct_graph(self):
            internal_graph = StaticTestComputation(one=1)
            self.connect_input_to_node("two", internal_graph, "two")
            self.set_output(internal_graph)


class StaticGraphWithElementaryNode(modelflow.graph.StaticGraph):
        two = modelflow.input.Input()
        one = modelflow.input.Input()

        def construct_graph(self):
            plus = modelflow.elementarynodes.Plus(arg1=1)
            self.connect_input_to_node("two", plus, "arg2")
            self.set_output(plus)


class TestStaticGraph(unittest.TestCase):
    def test_static_computation(self):
        static_graph = StaticTestComputation(one=1, two=2)
        self.assertEqual(3, static_graph())

    def test_static_graph_of_static_Graph(self):
        static_graph = StaticGraphofStaticGraph(two=2)
        self.assertEqual(3, static_graph())


class TestGraphPipeline(unittest.TestCase):

    class PipeTest(modelflow.graph.StaticGraph):
        in1 = modelflow.input.Input()
        in2 = modelflow.input.Input()
        in3 = modelflow.input.Input()

        def construct_graph(self):
            self.add_path(({"in1": "arg1",
                            "in2": "arg2"}, el.Plus()),
                          ({"in3": "arg2",
                            "%": "arg1"}, el.Plus()),
                          ({"in3": "enumerator",
                            "%": "denominator"}, el.Division()),
                          last_is_output=True)


    class PipeTestWithException(modelflow.graph.StaticGraph):
        in1 = modelflow.input.Input()
        in2 = modelflow.input.Input()
        in3 = modelflow.input.Input()

        def construct_graph(self):
            self.add_path(({"in1": "arg1",
                            "%": "arg2"}, el.Plus()),
                          ({"in3": "arg2",
                            "%": "arg1"}, el.Plus()),
                          ({"in3": "enumerator",
                            "%": "denominator"}, el.Division()),
                          last_is_output=True)

    def test_pipeline(self):
        pipe = self.PipeTest()
        self.assertEqual(4/(2+3+4), pipe(in1=2, in2=3, in3=4))

    def test_no_predecessor_input_exception(self):
        with self.assertRaises(Exception):
            pipe = self.PipeTestWithException()