from pyomo.contrib.appsi.base import PersistentBase
from pyomo.common.config import ConfigDict, ConfigValue, NonNegativeFloat, NonNegativeInt
from .cmodel import cmodel, cmodel_available
from typing import List, Optional
from pyomo.core.base.var import _GeneralVarData
from pyomo.core.base.param import _ParamData
from pyomo.core.base.constraint import _GeneralConstraintData
from pyomo.core.base.sos import _SOSConstraintData
from pyomo.core.base.objective import _GeneralObjectiveData, minimize, maximize
from pyomo.core.base.block import _BlockData
from pyomo.core.base import SymbolMap, TextLabeler
from pyomo.common.errors import InfeasibleConstraintException


class IntervalConfig(ConfigDict):
    """
    Attributes
    ----------
    feasibility_tol: float
    integer_tol: float
    improvement_tol: float
    max_iter: int
    """
    def __init__(self,
                 description=None,
                 doc=None,
                 implicit=False,
                 implicit_domain=None,
                 visibility=0):
        super(IntervalConfig, self).__init__(description=description,
                                           doc=doc,
                                           implicit=implicit,
                                           implicit_domain=implicit_domain,
                                           visibility=visibility)

        self.feasibility_tol: float = self.declare('feasibility_tol',
                                                   ConfigValue(domain=NonNegativeFloat, default=1e-8))
        self.integer_tol: float = self.declare('integer_tol',
                                               ConfigValue(domain=NonNegativeFloat, default=1e-5))
        self.improvement_tol: float = self.declare('improvement_tol',
                                                   ConfigValue(domain=NonNegativeFloat, default=1e-4))
        self.max_iter: int = self.declare('max_iter',
                                          ConfigValue(domain=NonNegativeInt, default=10))
        self.deactivate_satisfied_constraints: bool = self.declare('deactivate_satisfied_constraints',
                                                                   ConfigValue(domain=bool, default=False))


class IntervalTightener(PersistentBase):
    def __init__(self):
        super(IntervalTightener, self).__init__()
        self._config = IntervalConfig()
        self._cmodel = None
        self._var_map = dict()
        self._con_map = dict()
        self._param_map = dict()
        self._rvar_map = dict()
        self._rcon_map = dict()
        self._pyomo_expr_types = cmodel.PyomoExprTypes()
        self._symbolic_solver_labels: bool = False
        self._symbol_map = SymbolMap()
        self._var_labeler = None
        self._con_labeler = None
        self._param_labeler = None
        self._obj_labeler = None
        self._objective = None

    @property
    def config(self):
        return self._config

    @config.setter
    def config(self, val: IntervalConfig):
        self._config = val

    def set_instance(self, model, symbolic_solver_labels: Optional[bool] = None):
        saved_config = self.config
        saved_update_config = self.update_config

        self.__init__()
        self.config = saved_config
        self.update_config = saved_update_config
        self._expr_types = cmodel.PyomoExprTypes()

        if symbolic_solver_labels is not None:
            self._symbolic_solver_labels = symbolic_solver_labels
        if self._symbolic_solver_labels:
            self._var_labeler = TextLabeler()
            self._con_labeler = TextLabeler()
            self._param_labeler = TextLabeler()
            self._obj_labeler = TextLabeler()

        self._model = model
        self._cmodel = cmodel.FBBTModel()
        self.add_block(model)
        if self._objective is None:
            self.set_objective(None)

    def _add_variables(self, variables: List[_GeneralVarData]):
        if self._symbolic_solver_labels:
            set_name = True
            symbol_map = self._symbol_map
            labeler = self._var_labeler
        else:
            set_name = False
            symbol_map = None
            labeler = None
        cmodel.process_pyomo_vars(self._pyomo_expr_types, variables, self._var_map, self._param_map,
                                  self._vars, self._rvar_map, set_name, symbol_map, labeler, False)

    def _add_params(self, params: List[_ParamData]):
        cparams = cmodel.create_params(len(params))
        for ndx, p in enumerate(params):
            cp = cparams[ndx]
            cp.value = p.value
            self._param_map[id(p)] = cp
        if self._symbolic_solver_labels:
            for ndx, p in enumerate(params):
                cp = cparams[ndx]
                cp.name = self._symbol_map.getSymbol(p, self._param_labeler)

    def _add_constraints(self, cons: List[_GeneralConstraintData]):
        cmodel.process_fbbt_constraints(self._cmodel, self._pyomo_expr_types, cons, self._var_map, self._param_map,
                                        self._active_constraints, self._con_map, self._rcon_map)
        if self._symbolic_solver_labels:
            for c, cc in self._con_map.items():
                cc.name = self._symbol_map.getSymbol(c, self._con_labeler)

    def _add_sos_constraints(self, cons: List[_SOSConstraintData]):
        if len(cons) != 0:
            raise NotImplementedError('IntervalTightener does not support SOS constraints')

    def _remove_constraints(self, cons: List[_GeneralConstraintData]):
        if self._symbolic_solver_labels:
            for c in cons:
                self._symbol_map.removeSymbol(c)
                self._con_labeler.remove_obj(c)
        for c in cons:
            cc = self._con_map.pop(c)
            self._cmodel.remove_constraint(cc)
            del self._rcon_map[cc]

    def _remove_sos_constraints(self, cons: List[_SOSConstraintData]):
        if len(cons) != 0:
            raise NotImplementedError('IntervalTightener does not support SOS constraints')

    def _remove_variables(self, variables: List[_GeneralVarData]):
        if self._symbolic_solver_labels:
            for v in variables:
                self._symbol_map.removeSymbol(v)
                self._var_labeler.remove_obj(v)
        for v in variables:
            cvar = self._var_map.pop(id(v))
            del self._rvar_map[cvar]

    def _remove_params(self, params: List[_ParamData]):
        if self._symbolic_solver_labels:
            for p in params:
                self._symbol_map.removeSymbol(p)
                self._param_labeler.remove_obj(p)
        for p in params:
            del self._param_map[id(p)]

    def _update_variables(self, variables: List[_GeneralVarData]):
        cmodel.process_pyomo_vars(self._pyomo_expr_types, variables, self._var_map, self._param_map,
                                  self._vars, self._rvar_map, False, None, None, True)

    def update_params(self):
        for p_id, p in self._params.items():
            cp = self._param_map[p_id]
            cp.value = p.value

    def set_objective(self, obj: _GeneralObjectiveData):
        if self._symbolic_solver_labels:
            if self._objective is not None:
                self._symbol_map.removeSymbol(self._objective)
                self._obj_labeler.remove_obj(self._objective)
        super().set_objective(obj)

    def _set_objective(self, obj: _GeneralObjectiveData):
        if obj is None:
            ce = cmodel.Constant(0)
            sense = 0
        else:
            ce = cmodel.appsi_expr_from_pyomo_expr(obj.expr, self._var_map, self._param_map, self._pyomo_expr_types)
            if obj.sense is minimize:
                sense = 0
            else:
                sense = 1
        cobj = cmodel.FBBTObjective(ce)
        cobj.sense = sense
        self._cmodel.objective = cobj
        self._objective = obj
        if self._symbolic_solver_labels and obj is not None:
            cobj.name = self._symbol_map.getSymbol(obj, self._obj_labeler)

    def _update_pyomo_var_bounds(self):
        for cv, v in self._rvar_map.items():
            cv_lb = cv.get_lb()
            cv_ub = cv.get_ub()
            if -cmodel.inf < cv_lb:
                v.setlb(cv_lb)
                v_id = id(v)
                _v, _lb, _ub, _fixed, _domain, _value = self._vars[v_id]
                self._vars[v_id] = (_v, cv_lb, _ub, _fixed, _domain, _value)
            if cv_ub < cmodel.inf:
                v.setub(cv_ub)
                v_id = id(v)
                _v, _lb, _ub, _fixed, _domain, _value = self._vars[v_id]
                self._vars[v_id] = (_v, _lb, cv_ub, _fixed, _domain, _value)

    def _deactivate_satisfied_cons(self):
        cons_to_deactivate = list()
        if self.config.deactivate_satisfied_constraints:
            for c, cc in self._con_map.items():
                if not cc.active:
                    cons_to_deactivate.append(c)
        self.remove_constraints(cons_to_deactivate)
        for c in cons_to_deactivate:
            c.deactivate()

    def perform_fbbt(self, model: _BlockData, symbolic_solver_labels: Optional[bool] = None):
        if model is not self._model:
            self.set_instance(model, symbolic_solver_labels=symbolic_solver_labels)
        else:
            if symbolic_solver_labels is not None and symbolic_solver_labels != self._symbolic_solver_labels:
                raise RuntimeError('symbolic_solver_labels can only be changed through the set_instance method. '
                                   'Please either use set_instance or create a new instance of IntervalTightener.')
            self.update()
        try:
            n_iter = self._cmodel.perform_fbbt(self.config.feasibility_tol, self.config.integer_tol,
                                               self.config.improvement_tol, self.config.max_iter,
                                               self.config.deactivate_satisfied_constraints)
        finally:
            # we want to make sure the pyomo model and cmodel stay in sync
            # even if an exception is raised and caught
            self._update_pyomo_var_bounds()
            self._deactivate_satisfied_cons()
        return n_iter

    def perform_fbbt_with_seed(self, model: _BlockData, seed_var: _GeneralVarData):
        if model is not self._model:
            self.set_instance(model)
        else:
            self.update()
        try:
            n_iter = self._cmodel.perform_fbbt_with_seed(self._var_map[id(seed_var)], self.config.feasibility_tol,
                                                         self.config.integer_tol, self.config.improvement_tol,
                                                         self.config.max_iter,
                                                         self.config.deactivate_satisfied_constraints)
        finally:
            # we want to make sure the pyomo model and cmodel stay in sync
            # even if an exception is raised and caught
            self._update_pyomo_var_bounds()
            self._deactivate_satisfied_cons()
        return n_iter
