#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________

import collections

from pyomo.common.collections import ComponentMap
from pyomo.common.errors import MouseTrap
from pyomo.core.expr.expr_common import ExpressionType
from pyomo.core.expr.visitor import StreamBasedExpressionVisitor
from pyomo.core.expr.numeric_expr import NumericExpression
from pyomo.core.expr.relational_expr import RelationalExpression
import pyomo.core.expr.current as EXPR
import pyomo.core.expr.logical_expr as LE
from pyomo.core.base import (
    Binary, Constraint, ConstraintList, NonNegativeIntegers, VarList, value)
import pyomo.core.base.boolean_var as BV
from pyomo.core.base.expression import ScalarExpression, _GeneralExpressionData
from pyomo.core.base.param import ScalarParam, _ParamData
from pyomo.core.base.var import ScalarVar, _GeneralVarData
from pyomo.gdp.disjunct import AutoLinkedBooleanVar, Disjunct, Disjunction

def _dispatch_boolean_var(visitor, node):
    if node not in visitor.boolean_to_binary_map:
        binary = node.get_associated_binary()
        if binary is not None:
            visitor.boolean_to_binary_map[node] = binary
        else:
            z = visitor.z_vars.add()
            visitor.boolean_to_binary_map[node] = z
            node.associate_binary_var(z)
    return False, visitor.boolean_to_binary_map[node]

def _dispatch_var(visitor, node):
    return False, node

def _dispatch_param(visitor, node):
    if int(value(node)) == value(node):
        return False, node
    else:
        raise ValueError("Found non-integer valued Param '%s' in a logical "
                         "expression. This cannot be written to a disjunctive "
                         "form." % node.name)

def _dispatch_expression(visitor, node):
    return False, node.expr

def _before_relational_expr(visitor, node):
    raise MouseTrap("The RelationalExpression '%s' was used as a Boolean term "
                    "in a logical proposition. This is not yet supported "
                    "when transforming to disjunctive form." % node)

def _dispatch_not(visitor, node, a):
    # z == !a
    if a not in visitor.expansions:
        z = visitor.z_vars.add()
        visitor.constraints.add(z == 1 - a)
        visitor.expansions[a] = z
    return visitor.expansions[a]

def _dispatch_implication(visitor, node, a, b):
    # z == !a v b
    return _dispatch_or(visitor, node, 1 - a, b)

def _dispatch_equivalence(visitor, node, a, b):
    # z == (!a v b) ^ (a v !b)
    return _dispatch_and(
        visitor, node,
        _dispatch_or(visitor, node, 1 - a, b),
        _dispatch_or(visitor, node, a, 1 - b),
    )

def _dispatch_and(visitor, node, *args):
    # z == a ^ b ^ ...
    z = visitor.z_vars.add()
    for arg in args:
        visitor.constraints.add(arg >= z)
    return z

def _dispatch_or(visitor, node, *args):
    # z == a v b v ...
    # (!z v a v b v ...) ^ (z v !a) ^ (z v !b) ^ ...
    z = visitor.z_vars.add()
    visitor.constraints.add((1 - z) + sum(args) >= 1)
    for arg in args:
        visitor.constraints.add(z + (1 - arg) >= 1)
    return z

def _dispatch_xor(visitor, node, a, b):
    # z == a XOR b
    # This is a special case of exactly
    return _dispatch_exactly(visitor, node, 1, a, b)

def _get_integer_value(n, node):
    if n.__class__ in EXPR.native_numeric_types and int(n) == n:
        return n
    if n.__class__ not in EXPR.native_types:
        if n.is_potentially_variable():
            # [ESJ 11/22]: This is probably worth supporting sometime, but right
            # now we are abiding by what docplex allows in their 'count'
            # function. Part of supporting this will be making sure we catch
            # strict inequalities in the GDP transformations. Because if we
            # don't know that n is integer-valued we will be forced to write
            # strict inequalities instead of incrememting or decrementing by 1
            # in the disjunctions.
            raise MouseTrap(
                "The first argument '%s' to '%s' is potentially variable. "
                "This may be a mathematically coherent expression; However "
                "it is not yet supported to convert it to a disjunctive "
                "program." % (n, node))
        else:
            return n
    raise ValueError("The first argument to '%s' must be an integer.\n\t"
                     "Recieved: %s" % (node, n))

def _dispatch_exactly(visitor, node, *args):
    # z = sum(args[1:] == args[0]
    # This is currently implemented as:
    # [sum(args[1:] = n] v [[sum(args[1:]) < n] v [sum(args[1:]) > n]]
    M = len(args) - 1
    n = _get_integer_value(args[0], node)
    sum_expr = sum(args[1:])
    equality_disj = visitor.disjuncts[len(visitor.disjuncts)]
    equality_disj.constraint = Constraint(expr=sum_expr == n)
    inequality_disj = visitor.disjuncts[len(visitor.disjuncts)]
    inequality_disj.disjunction = Disjunction(
        expr=[[sum_expr <= n - 1], [sum_expr >= n + 1]])
    visitor.disjunctions[len(visitor.disjunctions)] = [equality_disj,
                                                       inequality_disj]
    return equality_disj.indicator_var.get_associated_binary()

def _dispatch_atleast(visitor, node, *args):
    # z = sum[args[1:] >= n
    # This is implemented as:
    # [sum(args[1:] >= n] v [sum(args[1:] < n]
    n = _get_integer_value(args[0], node)
    sum_expr = sum(args[1:])
    atleast_disj = visitor.disjuncts[len(visitor.disjuncts)]
    less_disj = visitor.disjuncts[len(visitor.disjuncts)]
    atleast_disj.constraint = Constraint(expr=sum_expr >= n)
    less_disj.constraint = Constraint(expr=sum_expr <= n - 1)
    visitor.disjunctions[len(visitor.disjunctions)] = [atleast_disj, less_disj]
    return atleast_disj.indicator_var.get_associated_binary()

def _dispatch_atmost(visitor, node, *args):
    # z = sum[args[1:] <= n
    # This is implemented as:
    # [sum(args[1:] <= n] v [sum(args[1:] > n]
    n = _get_integer_value(args[0], node)
    sum_expr = sum(args[1:])
    atmost_disj = visitor.disjuncts[len(visitor.disjuncts)]
    more_disj = visitor.disjuncts[len(visitor.disjuncts)]
    atmost_disj.constraint = Constraint(expr=sum_expr <= n)
    more_disj.constraint = Constraint(expr=sum_expr >= n + 1)
    visitor.disjunctions[len(visitor.disjunctions)] = [atmost_disj, more_disj]
    return atmost_disj.indicator_var.get_associated_binary()

_operator_dispatcher = {}
_operator_dispatcher[LE.ImplicationExpression] = _dispatch_implication
_operator_dispatcher[LE.EquivalenceExpression] = _dispatch_equivalence
_operator_dispatcher[LE.NotExpression] = _dispatch_not
_operator_dispatcher[LE.AndExpression] = _dispatch_and
_operator_dispatcher[LE.OrExpression] = _dispatch_or
_operator_dispatcher[LE.XorExpression] = _dispatch_xor
_operator_dispatcher[LE.ExactlyExpression] = _dispatch_exactly
_operator_dispatcher[LE.AtLeastExpression] = _dispatch_atleast
_operator_dispatcher[LE.AtMostExpression] = _dispatch_atmost

_before_child_dispatcher = {}
_before_child_dispatcher[BV.ScalarBooleanVar] = _dispatch_boolean_var
_before_child_dispatcher[BV._GeneralBooleanVarData] = _dispatch_boolean_var
_before_child_dispatcher[AutoLinkedBooleanVar] = _dispatch_boolean_var
_before_child_dispatcher[_ParamData] = _dispatch_param
_before_child_dispatcher[ScalarParam] = _dispatch_param
# for the moment, these are all just so we can get good error messages when we
# don't handle them:
_before_child_dispatcher[ScalarVar] = _dispatch_var
_before_child_dispatcher[_GeneralVarData] = _dispatch_var
_before_child_dispatcher[_GeneralExpressionData] = _dispatch_expression
_before_child_dispatcher[ScalarExpression] = _dispatch_expression

class LogicalToDisjunctiveVisitor(StreamBasedExpressionVisitor):
    """Converts BooleanExpressions to Linear (MIP) representation

    This converter eschews conjunctive normal form, and instead follows
    the well-trodden MINLP path of factorable programming.

    """

    def __init__(self):
        super().__init__()
        self.z_vars = VarList(domain=Binary)
        self.z_vars.construct()
        self.constraints = ConstraintList()
        self.disjuncts = Disjunct(NonNegativeIntegers, concrete=True)
        self.disjunctions = Disjunction(NonNegativeIntegers)
        self.disjunctions.construct()
        self.expansions = ComponentMap()
        self.boolean_to_binary_map = ComponentMap()

    def initializeWalker(self, expr):
        walk, result = self.beforeChild(None, expr, 0)
        if not walk:
            return False, self.finalizeResult(result)
        return True, expr

    def beforeChild(self, node, child, child_idx):
        if child.__class__ in EXPR.native_types:
            return False, child

        if child.is_numeric_type():
            # Just pass it through, we'll figure it out later
            return False, child
        if child.is_expression_type(ExpressionType.RELATIONAL):
            # Eventually we'll handle these. Right now we set a MouseTrap
            return _before_relational_expr(self, child)

        if not child.is_expression_type() or child.is_named_expression_type():
            return _before_child_dispatcher[child.__class__](self, child)

        return True, None

    def exitNode(self, node, data):
        return _operator_dispatcher[node.__class__](self, node, *data)

    def finalizeResult(self, result):
        # This LogicalExpression must evaluate to True (but note that we cannot
        # fix this variable to 1 since this logical expression could be living
        # on a Disjunct and later need to be relaxed.)
        self.constraints.add(result >= 1)
        return result
