#########################################################################
#
# Date: Aug 2004  Author: Daniel Stoffler
#
#       stoffler@scripps.edu
#
#       The Scripps Research Institute (TSRI)
#       Molecular Graphics Lab
#       La Jolla, CA 92037, USA
#
# Copyright: Daniel Stoffler, and TSRI
#
#########################################################################

import sys, os
from time import sleep

from NetworkEditor.macros import MacroNode

from Vision.VPE import VisualProgramingEnvironment
from Vision.StandardNodes import stdlib, DialNE, Range, Iterate, Pass, Print,\
     Accumulate

ed = None

###############################
## implement setUp and tearDown
###############################

def setUp():
    global ed
    ed = VisualProgramingEnvironment("test scheduling",
                                     withShell=0,
                                     visibleWidth=400, visibleHeight=300)
    ed.root.update()
    ed.configure(withThreads=0)
    ed.addLibraryInstance(stdlib, 'Vision.StandardNodes', 'stdlib')
    ed.root.update()
    

def tearDown():
    ed.exit_cb()
    import gc
    gc.collect()

##########################
## Helper methods
##########################

def pause(sleepTime=0.1):
    ed.master.update()
    sleep(sleepTime)


def test_01_RangeIterateAccumulator():
    """test if the Range, Iterator and Accumulator network work."""
    masterNet = ed.currentNetwork
    # add RANGE
    node0 = Range(constrkw = {}, name='range', library=stdlib)
    masterNet.addNode(node0,195,4)
    node0.inputPorts[1].widget.set(10,0)
    apply(node0.configure, (), {'expanded': True})
    # add ITERATE
    node1 = Iterate(constrkw = {}, name='iterate', library=stdlib)
    masterNet.addNode(node1,198,167)
    apply(node1.outputPorts[0].configure, (), {'datatype': 'int'})
    # add ACCUMULATOR
    node2 = Accumulate(constrkw = {}, name='accum', library=stdlib)
    masterNet.addNode(node2,216,272)
    apply(node2.inputPorts[0].configure, (), {'datatype': 'int'})
    # connect nodes
    masterNet.connectNodes(
        node0, node1, "data", "listToLoopOver", blocking=True)
    masterNet.connectNodes(
        node1, node2, "oneItem", "data", blocking=True)
    masterNet.connectNodes(
        node1, node2, "begin", "begin", blocking=True)
    masterNet.connectNodes(
        node1, node2, "end", "output", blocking=True)
    pause()
    # run network
    ed.currentNetwork.run()
    # test if this all worked
    data = node2.outputPorts[0].data
    assert data is not None, "Expected data, got %s"%data
    assert data == range(10), "Expected %s, got %s"%(range(10), data)


def test_02_IterateAccumRunMain():
    """test if we can run an iterate node inside a macro network, which is
    connected to an Accum node, which is connected to the MacroOutputNode,
    and then run the MAIN network. The Acucm node is used to check for the
    correct result"""
    masterNet = ed.currentNetwork
    # Create macro, add Iterate into macro, add accum node to macro,
    # connect iterate to accum, connect accum to Output Node
    node0 = MacroNode(name='macro0')
    masterNet.addNode(node0, 78, 176)
    node3 = Iterate(constrkw = {}, name='iterate', library=stdlib)
    node0.macroNetwork.addNode(node3,208,94)
    apply(node3.outputPorts[0].configure, (), {'datatype': 'int'})
    node4 = Accumulate(constrkw = {}, name='accum', library=stdlib)
    node0.macroNetwork.addNode(node4,218,187)
    apply(node4.inputPorts[0].configure, (), {'datatype': 'int'})
    node1 = node0.macroNetwork.ipNode
    node0.macroNetwork.connectNodes(
        node1, node3, "new", "listToLoopOver", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node4, "oneItem", "data", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node4, "begin", "begin", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node4, "end", "output", blocking=True)
    node2 = node0.macroNetwork.opNode
    node0.macroNetwork.connectNodes(
        node4, node2, "listOfValues", "new", blocking=True)
    node0.shrink()
    # next, add Range node
    node5 = Range(constrkw = {}, name='range', library=stdlib)
    masterNet.addNode(node5,188,18)
    node5.inputPorts[1].widget.set(5,0)
    apply(node5.configure, (), {'expanded': True})
    masterNet.connectNodes(
        node5, node0, "data", "iterate_listToLoopOver", blocking=True)
    pause()
    # now, run the MAIN network
    ed.currentNetwork.run()
    # check if the Accumulator node obtained data
    data = node4.outputPorts[0].data
    assert data is not None, "Expected data, got %s"%data
    # next, see if the iterate node produced the right data
    assert data == range(5), "Expected %s, got %s"%(range(5), data)
    # finally, check if the macro output node provided data
    assert node0.outputPorts[0].data is not None,\
           "Expected data, got %s"%node0.outputPorts[0].data
    assert node0.outputPorts[0].data == range(5),\
           "Expected %s, got %s"%(range(5), data)


def test_03_IterateAccumRunMacro():
    """test if we can run an iterate node inside a macro network, which is
    connected to an Accum node, which is connected to the MacroOutputNode,
    and then run the MACRO network. The Acucm node is used to check for the
    correct result"""
    masterNet = ed.currentNetwork
    masterNet.runOnNewData.value = True
    # Create macro, add Iterate into macro, add accum node to macro,
    # connect iterate to accum, connect accum to Output Node
    node0 = MacroNode(name='macro0')
    masterNet.addNode(node0, 78, 176)
    node3 = Iterate(constrkw = {}, name='iterate', library=stdlib)
    node0.macroNetwork.addNode(node3,208,94)
    apply(node3.outputPorts[0].configure, (), {'datatype': 'int'})
    node4 = Accumulate(constrkw = {}, name='accum', library=stdlib)
    node0.macroNetwork.addNode(node4,218,187)
    apply(node4.inputPorts[0].configure, (), {'datatype': 'int'})
    node1 = node0.macroNetwork.ipNode
    node0.macroNetwork.connectNodes(
        node1, node3, "new", "listToLoopOver", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node4, "oneItem", "data", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node4, "begin", "begin", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node4, "end", "output", blocking=True)
    node2 = node0.macroNetwork.opNode
    node0.macroNetwork.connectNodes(
        node4, node2, "listOfValues", "new", blocking=True)
    node0.shrink()
    # next, add Range node
    node5 = Range(constrkw = {}, name='range', library=stdlib)
    masterNet.addNode(node5,188,18)
    node5.inputPorts[1].widget.set(5, run=1)
    apply(node5.configure, (), {'expanded': True})
    masterNet.connectNodes(
        node5, node0, "data", "iterate_listToLoopOver", blocking=True)
    # go back to macro network
    node0.expand()
    pause()
    # ok, now run the macro network
    ed.currentNetwork.run()
    # check if the Accumulator node obtained data
    data = node4.outputPorts[0].data
    assert data is not None, "Expected data, got %s"%data
    # next, see if the iterate node produced the right data
    assert data == range(5), "Expected %s, got %s"%(range(5), data)
    # finally, check if the macro output node provided data
    assert node0.outputPorts[0].data is not None,\
           "Expected data, got %s"%node0.outputPorts[0].data 
    assert node0.outputPorts[0].data == range(5),\
           "Expected %s, got %s"%(range(5), node0.outputPorts[0].data)


def test_04_IterateRunMain():
    """test if we can run an iterate node inside a macro network, which is
    directly connected to the MacroOutputNode, and then run the MAIN network.
    We use the accumulator node in the main network to obtain the result of
    the iteration"""
    masterNet = ed.currentNetwork
    # Create macro, add Iterate into macro, connect to Input and Output Node
    node0 = MacroNode(name='macro0')
    masterNet.addNode(node0, 32, 154)
    node3 = Iterate(constrkw = {}, name='iterate', library=stdlib)
    node0.macroNetwork.addNode(node3,209,132)
    apply(node3.outputPorts[0].configure, (), {'datatype': 'int'})
    node1 = node0.macroNetwork.ipNode
    node0.macroNetwork.connectNodes(
        node1, node3, "new", "listToLoopOver", blocking=True)
    node2 = node0.macroNetwork.opNode
    node0.macroNetwork.connectNodes(
        node3, node2, "oneItem", "new", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node2, "begin", "new", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node2, "end", "new", blocking=True)
    node0.shrink()
    # adding node range
    node4 = Range(constrkw = {}, name='range', library=stdlib)
    masterNet.addNode(node4,170,12)
    node4.inputPorts[1].widget.set(5)
    apply(node4.configure, (), {'expanded': True})
    # node accum ##
    node5 = Accumulate(constrkw = {}, name='accum', library=stdlib)
    masterNet.addNode(node5,119,245)
    masterNet.connectNodes(
        node4, node0, "data", "iterate_listToLoopOver", blocking=True)
    masterNet.connectNodes(
        node0, node5, "iterate_oneItem", "data", blocking=True)
    masterNet.connectNodes(
        node0, node5, "iterate_begin", "begin", blocking=True)
    masterNet.connectNodes(
        node0, node5, "iterate_end", "output", blocking=True)
    pause()
    # now, hard run the main network
    ed.runCurrentNet_cb()
    # check if the Accumulator node obtained data
    data = node5.outputPorts[0].data
    assert data is not None, "Expected data, got %s"%data
    # next, see if the iterate node produced the right data
    assert data == range(5), "Expected %s, got %s"%(range(5), data)


def test_05_IterateRunMacro():
    """test if we can run an iterate node inside a macro network, which is
    directly connected to the MacroOutputNode, and then run the MACRO network.
    We use the accumulator node in the main network to obtain the result of
    the iteration"""
    masterNet = ed.currentNetwork
    # Create macro, add Iterate into macro, connect to Input and Output Node
    node0 = MacroNode(name='macro0')
    masterNet.addNode(node0, 32, 154)
    node3 = Iterate(constrkw = {}, name='iterate', library=stdlib)
    node0.macroNetwork.addNode(node3,209,132)
    apply(node3.outputPorts[0].configure, (), {'datatype': 'int'})
    node1 = node0.macroNetwork.ipNode
    node0.macroNetwork.connectNodes(
        node1, node3, "new", "listToLoopOver", blocking=True)
    node2 = node0.macroNetwork.opNode
    node0.macroNetwork.connectNodes(
        node3, node2, "oneItem", "new", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node2, "begin", "new", blocking=True)
    node0.macroNetwork.connectNodes(
        node3, node2, "end", "new", blocking=True)
    node0.shrink()
    # adding node range
    node4 = Range(constrkw = {}, name='range', library=stdlib)
    masterNet.addNode(node4,170,12)
    node4.inputPorts[1].widget.set(5)
    apply(node4.configure, (), {'expanded': True})
    # node accum ##
    node5 = Accumulate(constrkw = {}, name='accum', library=stdlib)
    masterNet.addNode(node5,119,245)
    masterNet.connectNodes(
        node4, node0, "data", "iterate_listToLoopOver", blocking=True)
    masterNet.connectNodes(
        node0, node5, "iterate_oneItem", "data", blocking=True)
    masterNet.connectNodes(
        node0, node5, "iterate_begin", "begin", blocking=True)
    masterNet.connectNodes(
        node0, node5, "iterate_end", "output", blocking=True)
    # now, run the MACRO network
    node0.expand()
    pause()

    ed.currentNetwork.run()
    node0.shrink()
    # now, hard run the main network
    ed.runCurrentNet_cb()
    # check if the Accumulator node obtained data
    data = node5.outputPorts[0].data
    assert data is not None, "Expected data, got %s"%data
    # next, see if the iterate node produced the right data
    assert data == range(5), "Expected %s, got %s"%(range(5), data)


def test_06_MacroWithRootNode():
    """test if a root node fires in a macro that is loaded from a file"""
    masterNet = ed.currentNetwork
    # Create macro, add dial node in macro and connect to output 
    node0 = MacroNode(name='macro0')
    masterNet.addNode(node0, 78, 176)
    node1 = DialNE()
    node0.macroNetwork.addNode(node1,208,94)
    node1.inputPorts[0].widget.set(1.00)
    node2 = node0.macroNetwork.opNode
    node0.macroNetwork.connectNodes(
        node1, node2, "value", "new", blocking=True)
    node0.shrink()

    # add print node in master network and connect macro output to print
    node3 =  Print(library =stdlib)
    masterNet.addNode(node3, 78, 226)

    # connect nodes
    masterNet.connectNodes(
        node0, node3, "Dial_value", "in1", blocking=True)

    # make sure there is no such file
    os.system("rm -f tmpDial_net.py*")
    # save the file
    ed.saveNetwork('tmpDial_net.py')
    # and delete the network
    ed.deleteNetwork(ed.currentNetwork)
    pause()

    # load the file
    ed.loadNetwork('tmpDial_net.py')
    # make sure there is no such file
    os.system("rm -f tmpDial_net.py*")
    pause()
    ed.currentNetwork.nodes[0].expand() 
    ed.networks['macro0'].run()
    ed.networks['tmpDial'].run()
    data=ed.networks['tmpDial'].nodes[1].inputPorts[0].data
    assert data == 1.0, "Expected data, got %s"%data
