from coopr.pyomo import *

# Problem borrowed from the lpsolve project documentation
# http://lpsolve.sourceforge.net/5.0/SOS.htm

# Don't forget that you can see that actual input to the solver if you use
# the -k command line flag:

# $ pyomo basic_sos2_example.py --solver=cplex -k

# A standalone SOS2 example, so the data comes from these functions instead of
# a .dat file.  Thus, define all the "cruft" up here, so it's clearer how the
# model is tied together down below.

def c_param_init ( model, v ):
    return (-1, -1, -3, -2, -2)[ v -1 ]  # -1 because Python is 0-based


def b_param_init ( model, c):
    return (30, 30)[ c -1 ]  # -1 because Python is 0-based


def A_param_init ( model, c, v):
    data = (
      (-1, -1,  1,  1,  0),
      ( 1,  0,  1, -3,  0),
    )

    return data[ c -1 ][ v -1 ]  # -1 because Python is 0-based


def obj_rule ( model ):
    objective_expression = sum(
      model.c[ i ] * model.x[ i ]
      for i in model.variable_set
    )

    return objective_expression


def constraint_rule ( model, c):
    constraint_equation = (
      model.b[ c ] >= sum(
        model.A[c, i] * model.x[ i ]
        for i in model.variable_set
      )
    )

    return constraint_equation


x1_constraint_rule = lambda M: M.x[1] <= 40
x2_constraint_rule = lambda M: M.x[2] <= 1
x5_constraint_rule = lambda M: M.x[5] <= 1


###############################################################################
# Cruft above complete, now we actually create and tie the model together

model = AbstractModel()
M = model

M.variable_set   = RangeSet(1, 5)
M.constraint_set = RangeSet(1, 2)

M.c = Param( M.variable_set, rule=c_param_init )

M.A = Param( M.constraint_set, M.variable_set, rule=A_param_init )
M.b = Param( M.constraint_set, rule=b_param_init )

M.x = Var( M.variable_set, within=PositiveReals )

M.obj = Objective( rule=obj_rule, sense=minimize )  # min "c transpose" X

# At first, this is little more than a standard form Ax=b ...
M.constraints   = Constraint( M.constraint_set, rule=constraint_rule )

# ... with a couple of extra constraints ...
M.x1_constraint = Constraint( rule=x1_constraint_rule )
M.x2_constraint = Constraint( rule=x2_constraint_rule )
M.x5_constraint = Constraint( rule=x5_constraint_rule )

# ... and finally, add the constraint for which this example was created
M.x_sos_vars  = SOSConstraint( var=M.x, sos=2 )
