from pyomo.core import *
from pyomo.pysp.scenariotree.tree_structure_model import \
    ScenarioTreeModelFromNetworkX
from pyomo.pysp.annotations import \
    (ConstraintStageAnnotation,
     StochasticConstraintBoundsAnnotation,
     StochasticConstraintBodyAnnotation,
     StochasticObjectiveAnnotation,
     StochasticVariableBoundsAnnotation)

def simple_twostage_scenario_tree():
    from pyomo.pysp.scenariotree.tree_structure_model \
        import CreateConcreteTwoStageScenarioTreeModel
    st_model = CreateConcreteTwoStageScenarioTreeModel(2)
    first_stage = st_model.Stages.first()
    second_stage = st_model.Stages.last()
    # First Stage
    st_model.StageCost[first_stage] = 'StageCost[1]'
    st_model.StageVariables[first_stage].add('x')
    # Second Stage
    st_model.StageCost[second_stage] = 'StageCost[2]'
    st_model.StageVariables[second_stage].add('y')
    return st_model

def simple_twostage_model():
    model = ConcreteModel()
    model.x = Var()
    model.y = Var()
    model.StageCost = Expression([1,2])
    model.StageCost.add(1, model.x)
    model.StageCost.add(2, model.y)
    model.c = ConstraintList()
    model.c.add(model.x >= 10)
    model.c.add(model.y >= 10)
    model.o = Objective(expr=sum_product(model.StageCost))
    return model

def simple_threestage_scenario_tree():
    from pyomo.pysp.scenariotree.tree_structure_model \
        import CreateConcreteTwoStageScenarioTreeModel
    import networkx
    G = networkx.balanced_tree(2,2,networkx.DiGraph())
    st_model = ScenarioTreeModelFromNetworkX(
        G,
        edge_probability_attribute=None)
    first_stage = st_model.Stages.first()
    second_stage = st_model.Stages.next(first_stage)
    third_stage = st_model.Stages.last()
    # First Stage
    st_model.StageCost[first_stage] = 'StageCost[1]'
    st_model.StageVariables[first_stage].add('x')
    # Second Stage
    st_model.StageCost[second_stage] = 'StageCost[2]'
    st_model.StageVariables[second_stage].add('y')
    # Third Stage
    st_model.StageCost[thrid_stage] = 'StageCost[3]'
    st_model.StageVariables[second_stage].add('z')
    return st_model

def simple_threestage_model():
    model = ConcreteModel()
    model.x = Var()
    model.y = Var()
    model.z = Var()
    model.StageCost = Expression([1,2,3])
    model.StageCost(1, model.x)
    model.StageCost(2, model.y)
    model.StageCost(3, model.z)
    model.c = ConstraintList()
    model.c.add(model.x >= 10)
    model.c.add(model.y >= 10)
    model.c.add(model.z >= 10)
    model.o = Objective(expr=sum_product(model.StageCost))
    return model
