from pyomo.core import *

def pysp_scenario_tree_model_callback():
    from pyomo.pysp.scenariotree.tree_structure_model \
        import CreateConcreteTwoStageScenarioTreeModel

    st_model = CreateConcreteTwoStageScenarioTreeModel(3)

    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')
    st_model.StageDerivedVariables[first_stage].add('y')

    # Second Stage
    st_model.StageCost[second_stage] = 'StageCost[2]'
    st_model.StageVariables[second_stage].add('z')
    st_model.StageDerivedVariables[second_stage].add('q')

    return st_model

cnt = 0
def pysp_instance_creation_callback(scenario_name, node_names):
    global cnt

    model = ConcreteModel()
    model.x = Var(bounds=(0,10))
    model.y = Expression(expr=model.x + 1)
    model.z = Var(bounds=(-10, 10))
    model.q = Expression(expr=model.z**2)
    model.StageCost = Expression([1,2])
    model.StageCost.add(1, model.x)
    model.StageCost.add(2, -model.z)
    model.o = Objective(expr=sum_product(model.StageCost))
    model.c = ConstraintList()
    model.c.add(model.x >= cnt)
    model.c.add(model.z <= cnt**2)

    cnt += 1

    return model
