#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and 
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain 
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________
#
# This file was originally part of PySP and Pyomo, available: https://github.com/Pyomo/pysp
# Copied with modification from pysp/tests/unit/test_scenariotree.py

import pyomo.common.unittest as unittest

from mpisppy.utils.pysp_model.tree_structure_model import \
    (ScenarioTreeModelFromNetworkX,
     CreateConcreteTwoStageScenarioTreeModel)
from mpisppy.utils.pysp_model.tree_structure import ScenarioTree
from mpisppy.utils.pysp_model.pysp_model import _get_nonant_list, _get_derived_nonant_list
from pyomo.core import (ConcreteModel,
                        Set,
                        Var,
                        Expression,
                        Objective,
                        Block,
                        value)
from pyomo.common.dependencies import (
    networkx, networkx_available as has_networkx
)

def _get_names(iterable):
    return [_.name for _ in iterable]

class TestScenarioTree(unittest.TestCase):

    def _get_block_model(self):
        model = ConcreteModel()
        model.s = Set(initialize=[1,2])
        b = Block(concrete=True)
        b.s = Set(initialize=[1,2])
        b.x = Var()
        b.X = Var(model.s)
        model.b1 = b.clone()
        model.b2 = b.clone()
        model.b3 = b.clone()
        model.b4 = b.clone()
        model.B1 = Block(model.s, rule=lambda _,i: b.clone())
        model.B2 = Block(model.s, rule=lambda _,i: b.clone())
        model.B3 = Block(model.s, rule=lambda _,i: b.clone())
        model.B4 = Block(model.s, rule=lambda _,i: b.clone())
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)
        return model

    def test_indexedblock_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("B1")
        st_model.StageDerivedVariables['Stage1'].add("B2")
        st_model.NodeVariables['RootNode'].add("B3")
        st_model.NodeDerivedVariables['RootNode'].add("B4")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 12
        assert len(root_derived_nonant_names) == 12

        for name in (
                     "B1[1].x", "B1[2].x",
                     "B3[1].x", "B3[2].x",
                     ):
            assert name in root_nonant_names

        for name in (
                     "B1[1].X", "B1[2].X",
                     "B3[1].X", "B3[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names

        for name in (
                     "B2[1].x", "B2[2].x",
                     "B4[1].x", "B4[2].x",
                     ):
            assert name in root_derived_nonant_names

        for name in (
                     "B2[1].X", "B2[2].X",
                     "B4[1].X", "B4[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names


    def test_indexedblock_wildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("B1[*]")
        st_model.StageDerivedVariables['Stage1'].add("B2[*]")
        st_model.NodeVariables['RootNode'].add("B3[*]")
        st_model.NodeDerivedVariables['RootNode'].add("B4[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 12
        assert len(root_derived_nonant_names) == 12

        for name in (
                     "B1[1].x", "B1[2].x",
                     "B3[1].x", "B3[2].x",
                     ):
            assert name in root_nonant_names

        for name in (
                     "B1[1].X", "B1[2].X",
                     "B3[1].X", "B3[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names

        for name in (
                     "B2[1].x", "B2[2].x",
                     "B4[1].x", "B4[2].x",
                     ):
            assert name in root_derived_nonant_names

        for name in (
                     "B2[1].X", "B2[2].X",
                     "B4[1].X", "B4[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names


    def test_singletonblock_wildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("b1[*]")
        st_model.StageDerivedVariables['Stage1'].add("b2[*]")
        st_model.NodeVariables['RootNode'].add("b3[*]")
        st_model.NodeDerivedVariables['RootNode'].add("b4[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("b1.x", "b3.x"):
            assert name in root_nonant_names
        for name in ("b1.X", "b3.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names
        for name in ("b2.x", "b4.x"):
            assert name in root_derived_nonant_names
        for name in ("b2.X", "b4.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names


    def test_singletonblock_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("b1")
        st_model.StageDerivedVariables['Stage1'].add("b2")
        st_model.NodeVariables['RootNode'].add("b3")
        st_model.NodeDerivedVariables['RootNode'].add("b4")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("b1.x", "b3.x"):
            assert name in root_nonant_names
        for name in ("b1.X", "b3.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names
        for name in ("b2.x", "b4.x"):
            assert name in root_derived_nonant_names
        for name in ("b2.X", "b4.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names

 
    def test_singletonvar_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x")
        st_model.StageDerivedVariables['Stage1'].add("y")
        st_model.NodeVariables['RootNode'].add("z")
        st_model.NodeDerivedVariables['RootNode'].add("q")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.x = Var()
        model.y = Var()
        model.z = Var()
        model.q = Var()
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 2 
        assert len(root_derived_nonant_names) == 2

        for name in ("x", "z"):
            assert name in root_nonant_names

        for name in ("y", "q"):
            assert name in root_derived_nonant_names


    def test_singletonvar_wildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x[*]")
        st_model.StageDerivedVariables['Stage1'].add("y[*]")
        st_model.NodeVariables['RootNode'].add("z[*]")
        st_model.NodeDerivedVariables['RootNode'].add("q[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.x = Var()
        model.y = Var()
        model.z = Var()
        model.q = Var()
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 2
        assert len(root_derived_nonant_names) == 2

        for name in ("x", "z"):
            assert name in root_nonant_names

        for name in ("y", "q"):
            assert name in root_derived_nonant_names


    def test_multiindexedvar_singlewildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x[*,* ]")
        st_model.StageDerivedVariables['Stage1'].add("y[ *,*]")
        st_model.NodeVariables['RootNode'].add("z[*,*]")
        st_model.NodeDerivedVariables['RootNode'].add("q[ * , * ]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.s = Set(initialize=[(1,'a'),(2,'b'),(3,'c')])
        model.x = Var(model.s)
        model.y = Var(model.s)
        model.z = Var(model.s)
        model.q = Var(model.s)
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("x", "z"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_nonant_names
        for name in ("y", "q"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_derived_nonant_names


    def test_indexedvar_indextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x[*]")
        st_model.StageDerivedVariables['Stage1'].add("y[*]")
        st_model.NodeVariables['RootNode'].add("z[*]")
        st_model.NodeDerivedVariables['RootNode'].add("q[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.s = Set(initialize=[1,2,3])
        model.x = Var(model.s)
        model.y = Var(model.s)
        model.z = Var(model.s)
        model.q = Var(model.s)
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("x", "z"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_nonant_names
        for name in ("y", "q"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_derived_nonant_names


    def test_indexedvar_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x")
        st_model.StageDerivedVariables['Stage1'].add("y")
        st_model.NodeVariables['RootNode'].add("z")
        st_model.NodeDerivedVariables['RootNode'].add("q")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.s = Set(initialize=[1,2,3])
        model.x = Var(model.s)
        model.y = Var(model.s)
        model.z = Var(model.s)
        model.q = Var(model.s)
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("x", "z"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_nonant_names
        for name in ("y", "q"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_derived_nonant_names


@unittest.skipIf(not has_networkx, "Requires networkx module")
class TestScenarioTreeFromNetworkX(unittest.TestCase):

    def test_empty(self):
        G = networkx.DiGraph()
        with self.assertRaises(networkx.NetworkXPointlessConcept):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_tree(self):
        G = networkx.DiGraph()
        G.add_node("1")
        G.add_node("2")
        G.add_edge("1", "2")
        G.add_edge("2", "1")
        with self.assertRaises(TypeError):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_directed(self):
        G = networkx.Graph()
        G.add_node("1")
        G.add_node("2")
        G.add_edge("1", "2")
        with self.assertRaises(TypeError):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_branching(self):
        G = networkx.DiGraph()
        G.add_node("1")
        G.add_node("2")
        G.add_node("R")
        G.add_edge("1", "R")
        G.add_edge("2", "R")
        with self.assertRaises(TypeError):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_enough_stages(self):
        G = networkx.DiGraph()
        G.add_node("R")
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(G)

    def test_missing_node_name(self):
        G = networkx.DiGraph()
        G.add_node("R", name="Root")
        G.add_node("C")
        G.add_edge("R", "C", weight=1)
        with self.assertRaises(KeyError):
            ScenarioTreeModelFromNetworkX(
                G,
                node_name_attribute="name")

    def test_missing_scenario_name(self):
        G = networkx.DiGraph()
        G.add_node("R", name="Root")
        G.add_node("C")
        G.add_edge("R", "C", weight=1)
        with self.assertRaises(KeyError):
            ScenarioTreeModelFromNetworkX(
                G,
                scenario_name_attribute="name")

    def test_missing_weight(self):
        G = networkx.DiGraph()
        G.add_node("R", name="Root")
        G.add_node("C", name="Child")
        G.add_edge("R", "C")
        with self.assertRaises(KeyError):
            ScenarioTreeModelFromNetworkX(G)

    def test_bad_weight1(self):
        G = networkx.DiGraph()
        G.add_node("R",)
        G.add_node("C",)
        G.add_edge("R", "C",weight=0.8)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(G)

    def test_bad_weight2(self):
        G = networkx.DiGraph()
        G.add_node("R")
        G.add_node("C1")
        G.add_edge("R", "C1", weight=0.8)
        G.add_node("C2")
        G.add_edge("R", "C2", weight=0.1)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(G)

    def test_bad_custom_stage_names1(self):
        G = networkx.DiGraph()
        G.add_node("R",)
        G.add_node("C1")
        G.add_edge("R", "C1", weight=1.0)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G, stage_names=["Stage1"])

    def test_bad_custom_stage_names2(self):
        G = networkx.DiGraph()
        G.add_node("R")
        G.add_node("C1")
        G.add_edge("R", "C1", weight=1.0)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G, stage_names=["Stage1","Stage1"])

    def test_two_stage(self):
        G = networkx.DiGraph()
        G.add_node("Root")
        G.add_node("Child1")
        G.add_edge("Root", "Child1", weight=0.8)
        G.add_node("Child2")
        G.add_edge("Root", "Child2", weight=0.2)
        model = ScenarioTreeModelFromNetworkX(G)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["Root", "Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Root"])),
            sorted(["Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Child1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["Child2"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["Child1", "Child2"]))
        self.assertEqual(value(model.ConditionalProbability["Root"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["Child1"]), 0.8)
        self.assertEqual(value(model.ConditionalProbability["Child2"]), 0.2)
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageVariables["Stage1"].add("x")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_two_stage_more_node_attributes(self):
        G = networkx.DiGraph()
        G.add_node("Root",
                   cost="c1",
                   variables=["x"],
                   derived_variables=["y"])
        G.add_node("Child1",
                   cost="c2",
                   variables=["q"],
                   derived_variables=["z"])
        G.add_edge("Root", "Child1", weight=0.8)
        G.add_node("Child2",
                   cost="c2",
                   variables=["q"],
                   derived_variables=["z"])
        G.add_edge("Root", "Child2", weight=0.2)
        model = ScenarioTreeModelFromNetworkX(G)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["Root", "Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Root"])),
            sorted(["Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Child1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["Child2"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["Child1", "Child2"]))
        self.assertEqual(value(model.ConditionalProbability["Root"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["Child1"]), 0.8)
        self.assertEqual(value(model.ConditionalProbability["Child2"]), 0.2)

        self.assertEqual(model.StageCost["Stage1"].value, None)
        self.assertEqual(list(model.StageVariables["Stage1"]), [])
        self.assertEqual(list(model.StageDerivedVariables["Stage1"]), [])

        self.assertEqual(model.NodeCost["Root"].value, "c1")
        self.assertEqual(list(model.NodeVariables["Root"]), ["x"])
        self.assertEqual(list(model.NodeDerivedVariables["Root"]), ["y"])

        self.assertEqual(model.StageCost["Stage2"].value, None)
        self.assertEqual(list(model.StageVariables["Stage2"]), [])
        self.assertEqual(list(model.StageDerivedVariables["Stage2"]), [])

        self.assertEqual(model.NodeCost["Child1"].value, "c2")
        self.assertEqual(list(model.NodeVariables["Child1"]), ["q"])
        self.assertEqual(list(model.NodeDerivedVariables["Child1"]), ["z"])

        self.assertEqual(model.NodeCost["Child2"].value, "c2")
        self.assertEqual(list(model.NodeVariables["Child2"]), ["q"])
        self.assertEqual(list(model.NodeDerivedVariables["Child2"]), ["z"])
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_two_stage_custom_names(self):
        G = networkx.DiGraph()
        G.add_node("R", label="Root")
        G.add_node("C1", label="Child1", scenario="S1")
        G.add_edge("R", "C1", probability=0.8)
        G.add_node("C2", label="Child2", scenario="S2")
        G.add_edge("R", "C2", probability=0.2)
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute="probability",
            node_name_attribute="label",
            stage_names=["T1","T2"],
            scenario_name_attribute="scenario")
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["T1", "T2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["Root", "Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Root"])),
            sorted(["Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Child1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["Child2"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["S1", "S2"]))
        self.assertEqual(value(model.ConditionalProbability["Root"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["Child1"]), 0.8)
        self.assertEqual(value(model.ConditionalProbability["Child2"]), 0.2)
        model.StageCost["T1"] = "c1"
        model.StageCost["T2"] = "c2"
        model.StageVariables["T1"].add("x")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_multi_stage(self):
        G = networkx.balanced_tree(3,2,networkx.DiGraph())
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2", "Stage3"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]))
        self.assertEqual(
            sorted(list(model.Children[0])),
            sorted([1,2,3]))
        self.assertEqual(
            sorted(list(model.Children[1])),
            sorted([4,5,6]))
        self.assertEqual(
            sorted(list(model.Children[2])),
            sorted([7,8,9]))
        self.assertEqual(
            sorted(list(model.Children[3])),
            sorted([10,11,12]))
        self.assertEqual(
            sorted(list(model.Children[4])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[5])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[6])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[7])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[8])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[9])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[10])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[11])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[12])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted([4, 5, 6, 7, 8, 9, 10, 11, 12]))
        self.assertEqual(value(model.ConditionalProbability[0]), 1.0)
        self.assertAlmostEqual(value(model.ConditionalProbability[1]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[2]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[3]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[4]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[5]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[6]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[7]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[8]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[9]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[10]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[11]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[12]), 1.0/3)
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageCost["Stage3"] = "c3"
        model.StageVariables["Stage1"].add("x")
        model.StageVariables["Stage2"].add("y")
        model.StageVariables["Stage3"].add("y")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_unbalanced(self):
        G = networkx.DiGraph()
        G.add_node("R")
        G.add_node("0")
        G.add_node("1")
        G.add_edge("R", "0")
        G.add_edge("R", "1")
        G.add_node("00")
        G.add_node("01")
        G.add_edge("0", "00")
        G.add_edge("0", "01")
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2", "Stage3"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["R","0","1","00","01"]))
        self.assertEqual(
            sorted(list(model.Children["R"])),
            sorted(["0", "1"]))
        self.assertEqual(
            sorted(list(model.Children["0"])),
            sorted(["00","01"]))
        self.assertEqual(
            sorted(list(model.Children["1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["00"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["01"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["00", "01", "1"]))
        self.assertEqual(value(model.ConditionalProbability["R"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["0"]), 0.5)
        self.assertEqual(value(model.ConditionalProbability["1"]), 0.5)
        self.assertEqual(value(model.ConditionalProbability["00"]), 0.5)
        self.assertEqual(value(model.ConditionalProbability["01"]), 0.5)
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageCost["Stage3"] = "c3"
        model.StageVariables["Stage1"].add("x")
        model.StageVariables["Stage2"].add("x")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_bundles(self):
        G = networkx.DiGraph()
        G.add_node("r")
        for i in range(4):
            G.add_node("u"+str(i), bundle=i%2)
            G.add_edge("r", "u"+str(i))
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["r", "u0", "u1", "u2", "u3"]))
        self.assertEqual(
            sorted(list(model.Children["r"])),
            sorted(["u0", "u1", "u2", "u3"]))
        for i in range(4):
            self.assertEqual(
                sorted(list(model.Children["u"+str(i)])),
                sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["u0", "u1", "u2", "u3"]))
        self.assertEqual(value(model.ConditionalProbability["r"]), 1.0)
        for i in range(4):
            self.assertEqual(value(model.ConditionalProbability["u"+str(i)]),
                             0.25)
        self.assertEqual(model.Bundling.value, True)
        self.assertEqual(list(model.Bundles), [0, 1])
        for k, bundle_name in enumerate(model.Bundles):
            self.assertEqual(list(model.BundleScenarios[bundle_name]),
                             ["u"+str(i) for i in range(4)
                              if i%2 == k])
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageVariables["Stage1"].add("x")
        ScenarioTree(scenariotreeinstance=model)

    def test_bundles_incomplete(self):
        G = networkx.DiGraph()
        G.add_node("r")
        for i in range(4):
            G.add_node("u"+str(i), bundle="B")
            G.add_edge("r", "u"+str(i))
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(model.Bundling.value, True)
        self.assertEqual(list(model.Bundles), ["B"])
        self.assertEqual(list(model.BundleScenarios["B"]),
                         ["u"+str(i) for i in range(4)])
        G.nodes["u0"]["bundle"] = None
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G,
                edge_probability_attribute=None)
        del G.nodes["u0"]["bundle"]
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G,
                edge_probability_attribute=None)


if __name__ == "__main__":
    unittest.main()

#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and 
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain 
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________
#
# This file was originally part of PySP and Pyomo, available: https://github.com/Pyomo/pysp
# Copied with modification from pysp/tests/unit/test_scenariotree.py

import pyomo.common.unittest as unittest

from mpisppy.utils.pysp_model.tree_structure_model import \
    (ScenarioTreeModelFromNetworkX,
     CreateConcreteTwoStageScenarioTreeModel)
from mpisppy.utils.pysp_model.tree_structure import ScenarioTree
from mpisppy.utils.pysp_model.pysp_model import _get_nonant_list, _get_derived_nonant_list
from pyomo.core import (ConcreteModel,
                        Set,
                        Var,
                        Expression,
                        Objective,
                        Block,
                        value)
from pyomo.common.dependencies import (
    networkx, networkx_available as has_networkx
)

def _get_names(iterable):
    return [_.name for _ in iterable]

class TestScenarioTree(unittest.TestCase):

    def _get_block_model(self):
        model = ConcreteModel()
        model.s = Set(initialize=[1,2])
        b = Block(concrete=True)
        b.s = Set(initialize=[1,2])
        b.x = Var()
        b.X = Var(model.s)
        model.b1 = b.clone()
        model.b2 = b.clone()
        model.b3 = b.clone()
        model.b4 = b.clone()
        model.B1 = Block(model.s, rule=lambda _,i: b.clone())
        model.B2 = Block(model.s, rule=lambda _,i: b.clone())
        model.B3 = Block(model.s, rule=lambda _,i: b.clone())
        model.B4 = Block(model.s, rule=lambda _,i: b.clone())
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)
        return model

    def test_indexedblock_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("B1")
        st_model.StageDerivedVariables['Stage1'].add("B2")
        st_model.NodeVariables['RootNode'].add("B3")
        st_model.NodeDerivedVariables['RootNode'].add("B4")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 12
        assert len(root_derived_nonant_names) == 12

        for name in (
                     "B1[1].x", "B1[2].x",
                     "B3[1].x", "B3[2].x",
                     ):
            assert name in root_nonant_names

        for name in (
                     "B1[1].X", "B1[2].X",
                     "B3[1].X", "B3[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names

        for name in (
                     "B2[1].x", "B2[2].x",
                     "B4[1].x", "B4[2].x",
                     ):
            assert name in root_derived_nonant_names

        for name in (
                     "B2[1].X", "B2[2].X",
                     "B4[1].X", "B4[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names


    def test_indexedblock_wildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("B1[*]")
        st_model.StageDerivedVariables['Stage1'].add("B2[*]")
        st_model.NodeVariables['RootNode'].add("B3[*]")
        st_model.NodeDerivedVariables['RootNode'].add("B4[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 12
        assert len(root_derived_nonant_names) == 12

        for name in (
                     "B1[1].x", "B1[2].x",
                     "B3[1].x", "B3[2].x",
                     ):
            assert name in root_nonant_names

        for name in (
                     "B1[1].X", "B1[2].X",
                     "B3[1].X", "B3[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names

        for name in (
                     "B2[1].x", "B2[2].x",
                     "B4[1].x", "B4[2].x",
                     ):
            assert name in root_derived_nonant_names

        for name in (
                     "B2[1].X", "B2[2].X",
                     "B4[1].X", "B4[2].X",
                     ):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names


    def test_singletonblock_wildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("b1[*]")
        st_model.StageDerivedVariables['Stage1'].add("b2[*]")
        st_model.NodeVariables['RootNode'].add("b3[*]")
        st_model.NodeDerivedVariables['RootNode'].add("b4[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("b1.x", "b3.x"):
            assert name in root_nonant_names
        for name in ("b1.X", "b3.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names
        for name in ("b2.x", "b4.x"):
            assert name in root_derived_nonant_names
        for name in ("b2.X", "b4.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names


    def test_singletonblock_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("b1")
        st_model.StageDerivedVariables['Stage1'].add("b2")
        st_model.NodeVariables['RootNode'].add("b3")
        st_model.NodeDerivedVariables['RootNode'].add("b4")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = self._get_block_model()

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("b1.x", "b3.x"):
            assert name in root_nonant_names
        for name in ("b1.X", "b3.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_nonant_names
        for name in ("b2.x", "b4.x"):
            assert name in root_derived_nonant_names
        for name in ("b2.X", "b4.X"):
            var = model.find_component(name)
            for vardata in var.values():
                assert vardata.name in root_derived_nonant_names

 
    def test_singletonvar_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x")
        st_model.StageDerivedVariables['Stage1'].add("y")
        st_model.NodeVariables['RootNode'].add("z")
        st_model.NodeDerivedVariables['RootNode'].add("q")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.x = Var()
        model.y = Var()
        model.z = Var()
        model.q = Var()
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 2 
        assert len(root_derived_nonant_names) == 2

        for name in ("x", "z"):
            assert name in root_nonant_names

        for name in ("y", "q"):
            assert name in root_derived_nonant_names


    def test_singletonvar_wildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x[*]")
        st_model.StageDerivedVariables['Stage1'].add("y[*]")
        st_model.NodeVariables['RootNode'].add("z[*]")
        st_model.NodeDerivedVariables['RootNode'].add("q[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.x = Var()
        model.y = Var()
        model.z = Var()
        model.q = Var()
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 2
        assert len(root_derived_nonant_names) == 2

        for name in ("x", "z"):
            assert name in root_nonant_names

        for name in ("y", "q"):
            assert name in root_derived_nonant_names


    def test_multiindexedvar_singlewildcardtemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x[*,* ]")
        st_model.StageDerivedVariables['Stage1'].add("y[ *,*]")
        st_model.NodeVariables['RootNode'].add("z[*,*]")
        st_model.NodeDerivedVariables['RootNode'].add("q[ * , * ]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.s = Set(initialize=[(1,'a'),(2,'b'),(3,'c')])
        model.x = Var(model.s)
        model.y = Var(model.s)
        model.z = Var(model.s)
        model.q = Var(model.s)
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("x", "z"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_nonant_names
        for name in ("y", "q"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_derived_nonant_names


    def test_indexedvar_indextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x[*]")
        st_model.StageDerivedVariables['Stage1'].add("y[*]")
        st_model.NodeVariables['RootNode'].add("z[*]")
        st_model.NodeDerivedVariables['RootNode'].add("q[*]")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.s = Set(initialize=[1,2,3])
        model.x = Var(model.s)
        model.y = Var(model.s)
        model.z = Var(model.s)
        model.q = Var(model.s)
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("x", "z"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_nonant_names
        for name in ("y", "q"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_derived_nonant_names


    def test_indexedvar_noindextemplate(self):
        st_model = CreateConcreteTwoStageScenarioTreeModel(1)
        st_model.StageVariables['Stage1'].add("x")
        st_model.StageDerivedVariables['Stage1'].add("y")
        st_model.NodeVariables['RootNode'].add("z")
        st_model.NodeDerivedVariables['RootNode'].add("q")
        st_model.StageCost['Stage1'] = "FirstStageCost"
        st_model.StageCost['Stage2'] = "SecondStageCost"

        scenario_tree = ScenarioTree(scenariotreeinstance=st_model)
        self.assertEqual(len(scenario_tree.stages), 2)
        self.assertEqual(len(scenario_tree.nodes), 2)
        self.assertEqual(len(scenario_tree.scenarios), 1)

        model = ConcreteModel()
        model.s = Set(initialize=[1,2,3])
        model.x = Var(model.s)
        model.y = Var(model.s)
        model.z = Var(model.s)
        model.q = Var(model.s)
        model.FirstStageCost = Expression(expr=0.0)
        model.SecondStageCost = Expression(expr=0.0)
        model.obj = Objective(expr=0.0)

        root = scenario_tree.findRootNode()
        root_nonant_names = _get_names(_get_nonant_list(model, root))
        root_derived_nonant_names = _get_names(_get_derived_nonant_list(model, root))
        assert len(root_nonant_names) == 6
        assert len(root_derived_nonant_names) == 6

        for name in ("x", "z"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_nonant_names
        for name in ("y", "q"):
            indexed_var = model.find_component(name)
            for index in model.s:
                var = indexed_var[index]
                assert var.name in root_derived_nonant_names


@unittest.skipIf(not has_networkx, "Requires networkx module")
class TestScenarioTreeFromNetworkX(unittest.TestCase):

    def test_empty(self):
        G = networkx.DiGraph()
        with self.assertRaises(networkx.NetworkXPointlessConcept):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_tree(self):
        G = networkx.DiGraph()
        G.add_node("1")
        G.add_node("2")
        G.add_edge("1", "2")
        G.add_edge("2", "1")
        with self.assertRaises(TypeError):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_directed(self):
        G = networkx.Graph()
        G.add_node("1")
        G.add_node("2")
        G.add_edge("1", "2")
        with self.assertRaises(TypeError):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_branching(self):
        G = networkx.DiGraph()
        G.add_node("1")
        G.add_node("2")
        G.add_node("R")
        G.add_edge("1", "R")
        G.add_edge("2", "R")
        with self.assertRaises(TypeError):
            ScenarioTreeModelFromNetworkX(G)

    def test_not_enough_stages(self):
        G = networkx.DiGraph()
        G.add_node("R")
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(G)

    def test_missing_node_name(self):
        G = networkx.DiGraph()
        G.add_node("R", name="Root")
        G.add_node("C")
        G.add_edge("R", "C", weight=1)
        with self.assertRaises(KeyError):
            ScenarioTreeModelFromNetworkX(
                G,
                node_name_attribute="name")

    def test_missing_scenario_name(self):
        G = networkx.DiGraph()
        G.add_node("R", name="Root")
        G.add_node("C")
        G.add_edge("R", "C", weight=1)
        with self.assertRaises(KeyError):
            ScenarioTreeModelFromNetworkX(
                G,
                scenario_name_attribute="name")

    def test_missing_weight(self):
        G = networkx.DiGraph()
        G.add_node("R", name="Root")
        G.add_node("C", name="Child")
        G.add_edge("R", "C")
        with self.assertRaises(KeyError):
            ScenarioTreeModelFromNetworkX(G)

    def test_bad_weight1(self):
        G = networkx.DiGraph()
        G.add_node("R",)
        G.add_node("C",)
        G.add_edge("R", "C",weight=0.8)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(G)

    def test_bad_weight2(self):
        G = networkx.DiGraph()
        G.add_node("R")
        G.add_node("C1")
        G.add_edge("R", "C1", weight=0.8)
        G.add_node("C2")
        G.add_edge("R", "C2", weight=0.1)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(G)

    def test_bad_custom_stage_names1(self):
        G = networkx.DiGraph()
        G.add_node("R",)
        G.add_node("C1")
        G.add_edge("R", "C1", weight=1.0)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G, stage_names=["Stage1"])

    def test_bad_custom_stage_names2(self):
        G = networkx.DiGraph()
        G.add_node("R")
        G.add_node("C1")
        G.add_edge("R", "C1", weight=1.0)
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G, stage_names=["Stage1","Stage1"])

    def test_two_stage(self):
        G = networkx.DiGraph()
        G.add_node("Root")
        G.add_node("Child1")
        G.add_edge("Root", "Child1", weight=0.8)
        G.add_node("Child2")
        G.add_edge("Root", "Child2", weight=0.2)
        model = ScenarioTreeModelFromNetworkX(G)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["Root", "Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Root"])),
            sorted(["Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Child1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["Child2"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["Child1", "Child2"]))
        self.assertEqual(value(model.ConditionalProbability["Root"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["Child1"]), 0.8)
        self.assertEqual(value(model.ConditionalProbability["Child2"]), 0.2)
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageVariables["Stage1"].add("x")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_two_stage_more_node_attributes(self):
        G = networkx.DiGraph()
        G.add_node("Root",
                   cost="c1",
                   variables=["x"],
                   derived_variables=["y"])
        G.add_node("Child1",
                   cost="c2",
                   variables=["q"],
                   derived_variables=["z"])
        G.add_edge("Root", "Child1", weight=0.8)
        G.add_node("Child2",
                   cost="c2",
                   variables=["q"],
                   derived_variables=["z"])
        G.add_edge("Root", "Child2", weight=0.2)
        model = ScenarioTreeModelFromNetworkX(G)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["Root", "Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Root"])),
            sorted(["Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Child1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["Child2"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["Child1", "Child2"]))
        self.assertEqual(value(model.ConditionalProbability["Root"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["Child1"]), 0.8)
        self.assertEqual(value(model.ConditionalProbability["Child2"]), 0.2)

        self.assertEqual(model.StageCost["Stage1"].value, None)
        self.assertEqual(list(model.StageVariables["Stage1"]), [])
        self.assertEqual(list(model.StageDerivedVariables["Stage1"]), [])

        self.assertEqual(model.NodeCost["Root"].value, "c1")
        self.assertEqual(list(model.NodeVariables["Root"]), ["x"])
        self.assertEqual(list(model.NodeDerivedVariables["Root"]), ["y"])

        self.assertEqual(model.StageCost["Stage2"].value, None)
        self.assertEqual(list(model.StageVariables["Stage2"]), [])
        self.assertEqual(list(model.StageDerivedVariables["Stage2"]), [])

        self.assertEqual(model.NodeCost["Child1"].value, "c2")
        self.assertEqual(list(model.NodeVariables["Child1"]), ["q"])
        self.assertEqual(list(model.NodeDerivedVariables["Child1"]), ["z"])

        self.assertEqual(model.NodeCost["Child2"].value, "c2")
        self.assertEqual(list(model.NodeVariables["Child2"]), ["q"])
        self.assertEqual(list(model.NodeDerivedVariables["Child2"]), ["z"])
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_two_stage_custom_names(self):
        G = networkx.DiGraph()
        G.add_node("R", label="Root")
        G.add_node("C1", label="Child1", scenario="S1")
        G.add_edge("R", "C1", probability=0.8)
        G.add_node("C2", label="Child2", scenario="S2")
        G.add_edge("R", "C2", probability=0.2)
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute="probability",
            node_name_attribute="label",
            stage_names=["T1","T2"],
            scenario_name_attribute="scenario")
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["T1", "T2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["Root", "Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Root"])),
            sorted(["Child1", "Child2"]))
        self.assertEqual(
            sorted(list(model.Children["Child1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["Child2"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["S1", "S2"]))
        self.assertEqual(value(model.ConditionalProbability["Root"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["Child1"]), 0.8)
        self.assertEqual(value(model.ConditionalProbability["Child2"]), 0.2)
        model.StageCost["T1"] = "c1"
        model.StageCost["T2"] = "c2"
        model.StageVariables["T1"].add("x")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_multi_stage(self):
        G = networkx.balanced_tree(3,2,networkx.DiGraph())
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2", "Stage3"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]))
        self.assertEqual(
            sorted(list(model.Children[0])),
            sorted([1,2,3]))
        self.assertEqual(
            sorted(list(model.Children[1])),
            sorted([4,5,6]))
        self.assertEqual(
            sorted(list(model.Children[2])),
            sorted([7,8,9]))
        self.assertEqual(
            sorted(list(model.Children[3])),
            sorted([10,11,12]))
        self.assertEqual(
            sorted(list(model.Children[4])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[5])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[6])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[7])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[8])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[9])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[10])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[11])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children[12])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted([4, 5, 6, 7, 8, 9, 10, 11, 12]))
        self.assertEqual(value(model.ConditionalProbability[0]), 1.0)
        self.assertAlmostEqual(value(model.ConditionalProbability[1]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[2]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[3]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[4]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[5]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[6]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[7]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[8]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[9]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[10]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[11]), 1.0/3)
        self.assertAlmostEqual(value(model.ConditionalProbability[12]), 1.0/3)
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageCost["Stage3"] = "c3"
        model.StageVariables["Stage1"].add("x")
        model.StageVariables["Stage2"].add("y")
        model.StageVariables["Stage3"].add("y")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_unbalanced(self):
        G = networkx.DiGraph()
        G.add_node("R")
        G.add_node("0")
        G.add_node("1")
        G.add_edge("R", "0")
        G.add_edge("R", "1")
        G.add_node("00")
        G.add_node("01")
        G.add_edge("0", "00")
        G.add_edge("0", "01")
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2", "Stage3"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["R","0","1","00","01"]))
        self.assertEqual(
            sorted(list(model.Children["R"])),
            sorted(["0", "1"]))
        self.assertEqual(
            sorted(list(model.Children["0"])),
            sorted(["00","01"]))
        self.assertEqual(
            sorted(list(model.Children["1"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["00"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Children["01"])),
            sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["00", "01", "1"]))
        self.assertEqual(value(model.ConditionalProbability["R"]), 1.0)
        self.assertEqual(value(model.ConditionalProbability["0"]), 0.5)
        self.assertEqual(value(model.ConditionalProbability["1"]), 0.5)
        self.assertEqual(value(model.ConditionalProbability["00"]), 0.5)
        self.assertEqual(value(model.ConditionalProbability["01"]), 0.5)
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageCost["Stage3"] = "c3"
        model.StageVariables["Stage1"].add("x")
        model.StageVariables["Stage2"].add("x")
        self.assertEqual(model.Bundling.value, False)
        self.assertEqual(list(model.Bundles), [])
        self.assertEqual(len(model.BundleScenarios), 0)
        ScenarioTree(scenariotreeinstance=model)

    def test_bundles(self):
        G = networkx.DiGraph()
        G.add_node("r")
        for i in range(4):
            G.add_node("u"+str(i), bundle=i%2)
            G.add_edge("r", "u"+str(i))
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(
            sorted(list(model.Stages)),
            sorted(["Stage1", "Stage2"]))
        self.assertEqual(
            sorted(list(model.Nodes)),
            sorted(["r", "u0", "u1", "u2", "u3"]))
        self.assertEqual(
            sorted(list(model.Children["r"])),
            sorted(["u0", "u1", "u2", "u3"]))
        for i in range(4):
            self.assertEqual(
                sorted(list(model.Children["u"+str(i)])),
                sorted([]))
        self.assertEqual(
            sorted(list(model.Scenarios)),
            sorted(["u0", "u1", "u2", "u3"]))
        self.assertEqual(value(model.ConditionalProbability["r"]), 1.0)
        for i in range(4):
            self.assertEqual(value(model.ConditionalProbability["u"+str(i)]),
                             0.25)
        self.assertEqual(model.Bundling.value, True)
        self.assertEqual(list(model.Bundles), [0, 1])
        for k, bundle_name in enumerate(model.Bundles):
            self.assertEqual(list(model.BundleScenarios[bundle_name]),
                             ["u"+str(i) for i in range(4)
                              if i%2 == k])
        model.StageCost["Stage1"] = "c1"
        model.StageCost["Stage2"] = "c2"
        model.StageVariables["Stage1"].add("x")
        ScenarioTree(scenariotreeinstance=model)

    def test_bundles_incomplete(self):
        G = networkx.DiGraph()
        G.add_node("r")
        for i in range(4):
            G.add_node("u"+str(i), bundle="B")
            G.add_edge("r", "u"+str(i))
        model = ScenarioTreeModelFromNetworkX(
            G,
            edge_probability_attribute=None)
        self.assertEqual(model.Bundling.value, True)
        self.assertEqual(list(model.Bundles), ["B"])
        self.assertEqual(list(model.BundleScenarios["B"]),
                         ["u"+str(i) for i in range(4)])
        G.nodes["u0"]["bundle"] = None
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G,
                edge_probability_attribute=None)
        del G.nodes["u0"]["bundle"]
        with self.assertRaises(ValueError):
            ScenarioTreeModelFromNetworkX(
                G,
                edge_probability_attribute=None)


if __name__ == "__main__":
    unittest.main()

#  ___________________________________________________________________________ # #  Pyomo: Python Optimization Modeling Objects #  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC #  Under the terms of Contract DE-NA0003525 with National Technology and #  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain #  rights in this software. #  This software is distributed under the 3-clause BSD License. #  ___________________________________________________________________________ # # This file was originally part of PySP and Pyomo, available: https://github.com/Pyomo/pysp # Copied with modification from pysp/tests/unit/test_scenariotree.py