"""Transformation from BooleanVar and LogicalConstraint to Binary and
Constraints."""
from pyomo.common.collections import ComponentMap
from pyomo.common.modeling import unique_component_name
from pyomo.common.config import ConfigBlock, ConfigValue
from pyomo.contrib.fbbt.fbbt import compute_bounds_on_expr
from pyomo.core import (TransformationFactory, BooleanVar, VarList, Binary,
                        LogicalConstraint, Block, ConstraintList, native_types,
                        BooleanVarList)
from pyomo.core.base.block import _BlockData
from pyomo.core.base.boolean_var import (
    _DeprecatedImplicitAssociatedBinaryVariable)
from pyomo.core.expr.cnf_walker import to_cnf
from pyomo.core.expr.current import (
    AndExpression, OrExpression,
    NotExpression, AtLeastExpression,
    AtMostExpression, ExactlyExpression,
    special_boolean_atom_types,
    EqualityExpression,
    InequalityExpression,
    RangedExpression,
)
from pyomo.core.expr.numvalue import native_logical_types, value
from pyomo.core.expr.visitor import StreamBasedExpressionVisitor
from pyomo.core.expr.current import identify_variables
from pyomo.core.plugins.transform.hierarchy import IsomorphicTransformation
from pyomo.core.util import target_list

@TransformationFactory.register(
    "core.logical_to_linear",
    doc="Convert logic to linear constraints")
class LogicalToLinear(IsomorphicTransformation):
    """
    Re-encode logical constraints as linear constraints,
    converting Boolean variables to binary.
    """
    CONFIG = ConfigBlock('core.logical_to_linear')
    CONFIG.declare('targets', ConfigValue(
        default=None,
        domain=target_list,
        description="target or list of targets that will be relaxed",
        doc="""
        This specifies the list of LogicalConstraints to transform, or the 
        list of Blocks or Disjuncts on which to transform all of the 
        LogicalConstraints. Note that if the transformation is done out
        of place, the list of targets should be attached to the model before it
        is cloned, and the list will specify the targets on the cloned
        instance.
        """
    ))

    def _apply_to(self, model, **kwds):
        config = self.CONFIG(kwds.pop('options', {}))
        config.set_value(kwds)
        targets = config.targets
        if targets is None:
            targets = (model, )

        new_var_lists = ComponentMap()
        transBlocks = {}
        for t in targets:
            # If the user promises that the target is Block-like, we will go
            # with it. Note, however, that they can only use targets for
            # this--when we go searching for stuff to transform we will only
            # look on Blocks. And yes, this means we are ignoring Disjuncts. We
            # are in fact ignoring all GDP components because this
            # transformation is a promise only to transform LogicalConstraints
            # and the relevant BooleanVars, not to create an algebraic
            # model. (We are making this decision largely because having this
            # transformation do anything to GDP stuff is an assumption on how
            # the GDP will be solved, and it would be wrong to assume that a GDP
            # will *necessarily* be solved as an algebraic model. The star
            # example of not doing so being GDPopt.)
            if t.ctype is Block or isinstance(t, _BlockData):
                self._transform_block(t, model, new_var_lists, transBlocks)
            elif t.ctype is LogicalConstraint:
                if t.is_indexed():
                    self._transform_constraint(t, new_var_lists, transBlocks)
                else:
                    self._transform_constraintData(t, new_var_lists,
                                                   transBlocks)
            else:
                raise RuntimeError("Target '%s' was not a Block, Disjunct, or"
                                   " LogicalConstraint. It was of type %s "
                                   "and can't be transformed." % (t.name,
                                                                  type(t)))

    def _transform_boolean_varData(self, bool_vardata, new_varlists):
        # This transformation tries to group the binaries it creates for indexed
        # BooleanVars onto the same VarList. This won't work across separate
        # calls to the transformation, but within one call it's fine. So we have
        # two cases: 1) either we have created a VarList for this
        # BooleanVarData's parent_component, but have yet to add its binary to
        # said list, or 2) we have neither the binary nor the VarList

        parent_component = bool_vardata.parent_component()
        new_varlist = new_varlists.get(parent_component)
        if new_varlist is None and \
           bool_vardata.get_associated_binary() is None:
            # Case 2) we have neither the VarList nor an associated binary
            parent_block = bool_vardata.parent_block()
            new_var_list_name = unique_component_name(
                parent_block,
                parent_component.local_name + '_asbinary')
            new_varlist = VarList(domain=Binary)
            setattr(parent_block, new_var_list_name, new_varlist)
            new_varlists[parent_component] = new_varlist

        if bool_vardata.get_associated_binary() is None:
            # Case 1) we already have a VarList, but need to create the
            # associated binary
            new_binary_vardata = new_varlist.add()
            bool_vardata.associate_binary_var(new_binary_vardata)
            if bool_vardata.value is not None:
                new_binary_vardata.value = int(bool_vardata.value)
            if bool_vardata.fixed:
                new_binary_vardata.fix()

    def _transform_constraint(self, constraint, new_varlists, transBlocks):
        for i in constraint.keys(ordered=True):
            self._transform_constraintData(constraint[i], new_varlists,
                                           transBlocks)
        constraint.deactivate()

    def _transform_block(self, target_block, model, new_varlists, transBlocks):
        _blocks = target_block.values() if target_block.is_indexed() else \
                  (target_block,)
        for block in _blocks:
            for logical_constraint in block.component_objects(
                    ctype=LogicalConstraint, active=True, descend_into=Block):
                self._transform_constraint(logical_constraint, new_varlists,
                                           transBlocks)

            # This can go away when we deprecate this transformation
            # transforming BooleanVars. This just marks the BooleanVars as
            # "seen" so that if someone asks for their binary var later, we can
            # create it on the fly and complain.
            for bool_vardata in block.component_data_objects(
                    BooleanVar, descend_into=Block):
                if bool_vardata._associated_binary is None:
                    bool_vardata._associated_binary = \
                                _DeprecatedImplicitAssociatedBinaryVariable(
                                    bool_vardata)

    def _transform_constraintData(self, logical_constraint, new_varlists,
                                  transBlocks):
        # first find all the relevant BooleanVars and associate a binary (if
        # they don't have one already)
        for bool_vardata in identify_variables(logical_constraint.expr):
            if bool_vardata.ctype is BooleanVar:
                self._transform_boolean_varData(bool_vardata, new_varlists)

        # now create a transformation block on the constraint's parent block (if
        # we don't have one already)
        parent_block = logical_constraint.parent_block()
        xfrm_block = transBlocks.get(parent_block)
        if xfrm_block is None:
            xfrm_block = self._create_transformation_block(parent_block)
            transBlocks[parent_block] = xfrm_block
        new_constrlist = xfrm_block.transformed_constraints
        new_boolvarlist = xfrm_block.augmented_vars
        new_varlist = xfrm_block.augmented_vars_asbinary

        old_boolvarlist_length = len(new_boolvarlist)

        indicator_map = ComponentMap()
        cnf_statements = to_cnf(logical_constraint.body, new_boolvarlist,
                                indicator_map)
        logical_constraint.deactivate()

        # Associate new Boolean vars to new binary variables
        num_new = len(new_boolvarlist) - old_boolvarlist_length
        list_o_vars = list(new_boolvarlist.values())
        if num_new:
            for bool_vardata in list_o_vars[-num_new:]:
                new_binary_vardata = new_varlist.add()
                bool_vardata.associate_binary_var(new_binary_vardata)

        # Add constraints associated with each CNF statement
        for cnf_statement in cnf_statements:
            for linear_constraint in _cnf_to_linear_constraint_list(
                    cnf_statement):
                new_constrlist.add(expr=linear_constraint)

        # Add bigM associated with special atoms
        # Note: this ad-hoc reformulation may be revisited for tightness in the
        # future.
        old_varlist_length = len(new_varlist)
        for indicator_var, special_atom in indicator_map.items():
            for linear_constraint in _cnf_to_linear_constraint_list(
                    special_atom,
                    indicator_var,
                    new_varlist):
                new_constrlist.add(expr=linear_constraint)

        # Previous step may have added auxiliary binaries. Associate augmented
        # Booleans to them.
        num_new = len(new_varlist) - old_varlist_length
        list_o_vars = list(new_varlist.values())
        if num_new:
            for binary_vardata in list_o_vars[-num_new:]:
                new_bool_vardata = new_boolvarlist.add()
                new_bool_vardata.associate_binary_var(binary_vardata)

    def _create_transformation_block(self, context):
        new_xfrm_block_name = unique_component_name(context, 'logic_to_linear')
        new_xfrm_block = Block(doc="Transformation objects for logic_to_linear")
        setattr(context, new_xfrm_block_name, new_xfrm_block)

        new_xfrm_block.transformed_constraints = ConstraintList()
        new_xfrm_block.augmented_vars = BooleanVarList()
        new_xfrm_block.augmented_vars_asbinary = VarList( domain=Binary)

        return new_xfrm_block

def update_boolean_vars_from_binary(model, integer_tolerance=1e-5):
    """Updates all Boolean variables based on the value of their linked binary
    variables."""
    for boolean_var in model.component_data_objects(BooleanVar,
                                                    descend_into=Block):
        binary_var = boolean_var.get_associated_binary()
        if binary_var is not None and binary_var.value is not None:
            if abs(binary_var.value - 1) <= integer_tolerance:
                boolean_var.value = True
            elif abs(binary_var.value) <= integer_tolerance:
                boolean_var.value = False
            else:
                raise ValueError("Binary variable has non-{0,1} value: "
                                 "%s = %s" % (binary_var.name,
                                              binary_var.value))
            boolean_var.stale = binary_var.stale

def _cnf_to_linear_constraint_list(cnf_expr, indicator_var=None,
                                   binary_varlist=None):
    # Screen for constants
    if type(cnf_expr) in native_types or cnf_expr.is_constant():
        if value(cnf_expr) is True:
            return []
        else:
            raise ValueError(
                "Cannot build linear constraint for logical expression with "
                "constant value False: %s"
                % cnf_expr)
    if cnf_expr.is_expression_type():
        return CnfToLinearVisitor(indicator_var, binary_varlist).\
            walk_expression(cnf_expr)
    else:
        return [cnf_expr.get_associated_binary() == 1]  # Assume that cnf_expr
                                                        # is a BooleanVar

_numeric_relational_types = {InequalityExpression, EqualityExpression,
                             RangedExpression}

class CnfToLinearVisitor(StreamBasedExpressionVisitor):
    """Convert CNF logical constraint to linear constraints.

    Expected expression node types: AndExpression, OrExpression, NotExpression,
    AtLeastExpression, AtMostExpression, ExactlyExpression, _BooleanVarData

    """
    def __init__(self, indicator_var, binary_varlist):
        super(CnfToLinearVisitor, self).__init__()
        self._indicator = indicator_var
        self._binary_varlist = binary_varlist

    def exitNode(self, node, values):
        if type(node) == AndExpression:
            return list((v if type(v) in _numeric_relational_types else v == 1)
                        for v in values)
        elif type(node) == OrExpression:
            return sum(values) >= 1
        elif type(node) == NotExpression:
            return 1 - values[0]
        # Note: the following special atoms should only be encountered as root
        # nodes.  If they are encountered otherwise, something went wrong.
        sum_values = sum(values[1:])
        num_args = node.nargs() - 1  # number of logical arguments
        if self._indicator is None:
            if type(node) == AtLeastExpression:
                return sum_values >= values[0]
            elif type(node) == AtMostExpression:
                return sum_values <= values[0]
            elif type(node) == ExactlyExpression:
                return sum_values == values[0]
        else:
            rhs_lb, rhs_ub = compute_bounds_on_expr(values[0])
            if rhs_lb == float('-inf') or rhs_ub == float('inf'):
                raise ValueError( "Cannnot generate linear constraints for %s"
                                  "([N, *logical_args]) with unbounded N. "
                                  "Detected %s <= N <= %s." %
                                  (type(node).__name__, rhs_lb, rhs_ub) )
            indicator_binary = self._indicator.get_associated_binary()
            if type(node) == AtLeastExpression:
                return [
                    sum_values >= values[0] - rhs_ub * (1 - indicator_binary),
                    sum_values <= values[0] - 1 + (-(rhs_lb - 1) + num_args) * \
                    indicator_binary
                ]
            elif type(node) == AtMostExpression:
                return [
                    sum_values <= values[0] + (-rhs_lb + num_args) * \
                    (1 - indicator_binary),
                    sum_values >= (values[0] + 1) - (rhs_ub + 1) * \
                    indicator_binary
                ]
            elif type(node) == ExactlyExpression:
                less_than_binary = self._binary_varlist.add()
                more_than_binary = self._binary_varlist.add()
                return [
                    sum_values <= values[0] + (-rhs_lb + num_args) * \
                    (1 - indicator_binary),
                    sum_values >= values[0] - rhs_ub * (1 - indicator_binary),
                    indicator_binary + less_than_binary + more_than_binary >= 1,
                    sum_values <= values[0] - 1 + (-(rhs_lb - 1) + num_args) * \
                    (1 - less_than_binary),
                    sum_values >= values[0] + 1 - (rhs_ub + 1) * \
                    (1 - more_than_binary),
                ]
            pass

    def beforeChild(self, node, child, child_idx):
        if type(node) in special_boolean_atom_types and child is node.args[0]:
            return False, child
        if type(child) in native_logical_types:
            return False, int(child)
        if type(child) in native_types:
            return False, child

        if child.is_expression_type():
            return True, None

        # Only thing left should be _BooleanVarData
        return False, child.get_associated_binary()

    def finalizeResult(self, result):
        if type(result) is list:
            return result
        elif type(result) in _numeric_relational_types:
            return [result]
        else:
            return [result == 1]
