1from typing import List
2from pyomo.core.base.param import _ParamData
3from pyomo.core.base.var import _GeneralVarData
4from pyomo.core.base.constraint import _GeneralConstraintData
5from pyomo.core.base.objective import _GeneralObjectiveData
6from pyomo.core.base.sos import _SOSConstraintData
7from pyomo.core.base.block import _BlockData
8from pyomo.repn.standard_repn import generate_standard_repn
9from pyomo.core.expr.numvalue import value
10from pyomo.contrib.appsi.base import PersistentBase
11from pyomo.core.base import SymbolMap, NumericLabeler, TextLabeler
12from pyomo.common.timing import HierarchicalTimer
13from pyomo.core.kernel.objective import minimize, maximize
14from .config import WriterConfig
15from .cmodel_converter import PyomoToCModelWalker
16from ..cmodel import cmodel, cmodel_available
17
18id = id
19
20
21class LPWriter(PersistentBase):
22    def __init__(self):
23        super(LPWriter, self).__init__()
24        self._config = WriterConfig()
25        self._writer = None
26        self._symbol_map = SymbolMap()
27        self._var_labeler = None
28        self._con_labeler = None
29        self._param_labeler = None
30        self._obj_labeler = None
31        self._pyomo_var_to_solver_var_map = dict()
32        self._pyomo_con_to_solver_con_map = dict()
33        self._solver_var_to_pyomo_var_map = dict()
34        self._solver_con_to_pyomo_con_map = dict()
35        self._pyomo_param_to_solver_param_map = dict()
36        self._walker = PyomoToCModelWalker(self._pyomo_var_to_solver_var_map, self._pyomo_param_to_solver_param_map)
37
38    @property
39    def config(self):
40        return self._config
41
42    @config.setter
43    def config(self, val: WriterConfig):
44        self._config = val
45
46    def set_instance(self, model):
47        saved_config = self.config
48        saved_update_config = self.update_config
49        self.__init__()
50        self.config = saved_config
51        self.update_config = saved_update_config
52        self._model = model
53
54        if self.config.symbolic_solver_labels:
55            self._var_labeler = TextLabeler()
56            self._con_labeler = TextLabeler()
57            self._param_labeler = TextLabeler()
58            self._obj_labeler = TextLabeler()
59        else:
60            self._var_labeler = NumericLabeler('x')
61            self._con_labeler = NumericLabeler('c')
62            self._param_labeler = NumericLabeler('p')
63            self._obj_labeler = NumericLabeler('obj')
64
65        self._writer = cmodel.LPWriter()
66
67        self.add_block(model)
68        if self._objective is None:
69            self.set_objective(None)
70
71    def _add_variables(self, variables: List[_GeneralVarData]):
72        cvars = cmodel.create_vars(len(variables))
73        for ndx, v in enumerate(variables):
74            cv = cvars[ndx]
75            cv.name = self._symbol_map.getSymbol(v, self._var_labeler)
76            if v.is_binary():
77                cv.domain = 'binary'
78            elif v.is_integer():
79                cv.domain = 'integer'
80            else:
81                assert v.is_continuous(), 'LP writer only supports continuous, binary, and integer variables'
82                cv.domain = 'continuous'
83            _, lb, ub, v_is_fixed, v_domain, v_value = self._vars[id(v)]
84            if lb is not None:
85                cv.lb = lb
86            if ub is not None:
87                cv.ub = ub
88            if v_value is not None:
89                cv.value = v_value
90            if v_is_fixed:
91                cv.fixed = True
92            self._pyomo_var_to_solver_var_map[id(v)] = cv
93            self._solver_var_to_pyomo_var_map[cv] = v
94
95    def _add_params(self, params: List[_ParamData]):
96        cparams = cmodel.create_params(len(params))
97        for ndx, p in enumerate(params):
98            cp = cparams[ndx]
99            cp.name = self._symbol_map.getSymbol(p, self._param_labeler)
100            cp.value = p.value
101            self._pyomo_param_to_solver_param_map[id(p)] = cp
102
103    def _add_constraints(self, cons: List[_GeneralConstraintData]):
104        cmodel.process_lp_constraints(cons, self)
105
106    def _add_sos_constraints(self, cons: List[_SOSConstraintData]):
107        if len(cons) != 0:
108            raise NotImplementedError('LP writer does not yet support SOS constraints')
109
110    def _remove_constraints(self, cons: List[_GeneralConstraintData]):
111        for c in cons:
112            cc = self._pyomo_con_to_solver_con_map.pop(c)
113            self._writer.remove_constraint(cc)
114            self._symbol_map.removeSymbol(c)
115            self._con_labeler.remove_obj(c)
116            del self._solver_con_to_pyomo_con_map[cc]
117
118    def _remove_sos_constraints(self, cons: List[_SOSConstraintData]):
119        if len(cons) != 0:
120            raise NotImplementedError('LP writer does not yet support SOS constraints')
121
122    def _remove_variables(self, variables: List[_GeneralVarData]):
123        for v in variables:
124            cvar = self._pyomo_var_to_solver_var_map.pop(id(v))
125            del self._solver_var_to_pyomo_var_map[cvar]
126            self._symbol_map.removeSymbol(v)
127            self._var_labeler.remove_obj(v)
128
129    def _remove_params(self, params: List[_ParamData]):
130        for p in params:
131            del self._pyomo_param_to_solver_param_map[id(p)]
132            self._symbol_map.removeSymbol(p)
133            self._param_labeler.remove_obj(p)
134
135    def _update_variables(self, variables: List[_GeneralVarData]):
136        for v in variables:
137            cv = self._pyomo_var_to_solver_var_map[id(v)]
138            if v.is_binary():
139                cv.domain = 'binary'
140            elif v.is_integer():
141                cv.domain = 'integer'
142            else:
143                assert v.is_continuous(), 'LP writer only supports continuous, binary, and integer variables'
144                cv.domain = 'continuous'
145            lb = value(v.lb)
146            ub = value(v.ub)
147            if lb is None:
148                cv.lb = -cmodel.inf
149            else:
150                cv.lb = lb
151            if ub is None:
152                cv.ub = cmodel.inf
153            else:
154                cv.ub = ub
155            if v.value is not None:
156                cv.value = v.value
157            if v.is_fixed():
158                cv.fixed = True
159            else:
160                cv.fixed = False
161
162    def update_params(self):
163        for p_id, p in self._params.items():
164            cp = self._pyomo_param_to_solver_param_map[p_id]
165            cp.value = p.value
166
167    def _set_objective(self, obj: _GeneralObjectiveData):
168        if obj is None:
169            const = cmodel.Constant(0)
170            lin_coef = list()
171            lin_vars = list()
172            quad_coef = list()
173            quad_vars_1 = list()
174            quad_vars_2 = list()
175            sense = 0
176        else:
177            repn = generate_standard_repn(obj.expr, compute_values=False, quadratic=True)
178            const = self._walker.dfs_postorder_stack(repn.constant)
179            lin_coef = [self._walker.dfs_postorder_stack(i) for i in repn.linear_coefs]
180            lin_vars = [self._pyomo_var_to_solver_var_map[id(i)] for i in repn.linear_vars]
181            quad_coef = [self._walker.dfs_postorder_stack(i) for i in repn.quadratic_coefs]
182            quad_vars_1 = [self._pyomo_var_to_solver_var_map[id(i[0])] for i in repn.quadratic_vars]
183            quad_vars_2 = [self._pyomo_var_to_solver_var_map[id(i[1])] for i in repn.quadratic_vars]
184            if obj.sense is minimize:
185                sense = 0
186            else:
187                sense = 1
188        cobj = cmodel.LPObjective(const, lin_coef, lin_vars, quad_coef, quad_vars_1, quad_vars_2)
189        cobj.sense = sense
190        if obj is None:
191            cname = 'objective'
192        else:
193            cname = self._symbol_map.getSymbol(obj, self._obj_labeler)
194        cobj.name = cname
195        self._writer.objective = cobj
196
197    def write(self, model: _BlockData, filename: str, timer: HierarchicalTimer = None):
198        if timer is None:
199            timer = HierarchicalTimer()
200        if model is not self._model:
201            timer.start('set_instance')
202            self.set_instance(model)
203            timer.stop('set_instance')
204        else:
205            timer.start('update')
206            self.update(timer=timer)
207            timer.stop('update')
208        timer.start('write file')
209        self._writer.write(filename)
210        timer.stop('write file')
211
212    def get_vars(self):
213        return [self._solver_var_to_pyomo_var_map[i] for i in self._writer.get_solve_vars()]
214
215    def get_ordered_cons(self):
216        return [self._solver_con_to_pyomo_con_map[i] for i in self._writer.get_solve_cons()]
217
218    def get_active_objective(self):
219        return self._objective
220
221    @property
222    def symbol_map(self):
223        return self._symbol_map
224