import unittest

import modelflow.elementarynodes
import modelflow.graph


class TestDynamicGraph(unittest.TestCase):
    def test_inputs_added_to_channels(self):
        graph = modelflow.graph.DynamicGraph()
        input_list = ["one", "two", "three"]
        graph.add_inputs(["one", "two", "three"])
        self.assertEqual(input_list, graph.input_names)

    def test_get_input(self):
        graph = modelflow.graph.DynamicGraph()
        graph.add_inputs(['arg1'])
        graph['arg1'] = 2
        self.assertEqual(2, graph['arg1'])

    def test_computation(self):
        minus = modelflow.elementarynodes.Minus()
        div = modelflow.elementarynodes.Division()
        graph = modelflow.graph.DynamicGraph()
        graph.add_inputs(["one", "two", "three"])
        graph.connect_input_to_node("one", minus, "plus")
        graph.connect_input_to_node("two", minus, "minus")
        graph.connect_input_to_node("three", div, "denominator")
        graph.connect_node_to_node(minus, div, "enumerator")
        graph.set_output(div)
        self.assertEqual((1-2)/3, graph({"one": 1, "two": 2, "three": 3}))


class TestDynamicGraphOfDynamicGraphs(unittest.TestCase):
    def setUp(self):
        plus = modelflow.elementarynodes.Plus()
        div = modelflow.elementarynodes.Division()
        graph = modelflow.graph.DynamicGraph()
        graph.add_inputs(["one", "two", "three"])
        graph.connect_input_to_node("one", plus, "arg1")
        graph.connect_input_to_node("two", plus, "arg2")
        graph.connect_input_to_node("three", div, "denominator")
        graph.connect_node_to_node(plus, div, "enumerator")
        graph.set_output(div)
        graph['one'] = 1
        graph['three'] = 2
        # graph calculates (1 + two) / 2
        self.graph = graph


        graph2 = modelflow.graph.DynamicGraph()
        plus2 = modelflow.elementarynodes.Plus()
        graph2.add_inputs(['arg1', 'arg2'])
        graph2.connect_input_to_node('arg1', plus2, 'arg1')
        graph2.connect_input_to_node('arg2', plus2, 'arg2')
        graph2.set_output(plus2)
        # graph 2 calculates (arg1 + arg2)
        self.graph2 = graph2

    def test_first_graph(self):
        two = 2
        self.assertEqual((1+two)/2, self.graph(two=two))

    def test_second_graph(self):
        arg = 2
        self.assertEqual(arg+arg, self.graph2(arg1=arg, arg2=arg))

    def test_graph_of_graphs(self):
        graph, graph2 = self.graph, self.graph2

        graph_of_graphs = modelflow.graph.DynamicGraph()
        graph_of_graphs.add_inputs(['arg1'])
        graph_of_graphs.connect_input_to_node('arg1', graph2, 'arg1'),
        graph_of_graphs.connect_input_to_node('arg1', graph2, 'arg2'),
        graph_of_graphs.connect_node_to_node(graph2, graph, 'two')
        graph_of_graphs.set_output(graph)
        # graph_of_graphs calculates (1 + (arg1+arg1)) / 2

        arg1 = 2
        expected = (1 + (arg1+arg1)) / 2
        self.assertEqual(expected, graph_of_graphs(arg1=arg1))