#  _________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2014 Sandia Corporation.
#  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
#  the U.S. Government retains certain rights in this software.
#  This software is distributed under the BSD License.
#  _________________________________________________________________________

import six
import logging

from pyomo.core.base import Block, VarList, ConstraintList, Objective, Var, Constraint, maximize, ComponentUID
from pyomo.repn import generate_canonical_repn
from pyomo.repn.collect import collect_linear_terms
from pyomo.util.plugin import alias
from pyomo.mpec import ComplementarityList, complements
from pyomo.bilevel.plugins.transform import Base_BilevelTransformation
from pyomo.bilevel.components import SubModel


logger = logging.getLogger('pyomo.core')


class LinearComplementarity_BilevelTransformation(Base_BilevelTransformation):

    alias('bilevel.linear_mpec', doc="Generate a linear MPEC from the optimality conditions of the submodel")

    def __init__(self):
        super(LinearComplementarity_BilevelTransformation, self).__init__()

    def _apply_to(self, instance, **kwds):
        self._deterministic = kwds.get('deterministic',False)
        #
        # Process options
        #
        submodel = self._preprocess('bilevel.linear_mpec', instance, **kwds)
        instance.reclassify_component_type(submodel, Block)
        #
        # Create a block with optimality conditions
        #
        setattr(instance, self._submodel+'_kkt', self._add_optimality_conditions(instance, submodel))
        instance._transformation_data['bilevel.linear_mpec'].submodel_cuid = ComponentUID(submodel)
        instance._transformation_data['bilevel.linear_mpec'].block_cuid = ComponentUID(getattr(instance,self._submodel+'_kkt'))
        #-------------------------------------------------------------------------------
        #
        # Disable the original submodel and
        # execute the preprocessor
        #
        instance.reclassify_component_type(submodel, SubModel)
        submodel.deactivate()

    def _add_optimality_conditions(self, instance, submodel):
        """
        Add optimality conditions for the submodel

        This assumes that the original model has the form:

            min c1*x + d1*y
                A3*x <= b3
                A1*x + B1*y <= b1
                min c2*x + d2*y
                    y >= 0
                    A2*x + B2*y <= b2

        NOTE THE VARIABLE BOUNDS!
        """ 
        #
        # Populate the block with the linear constraints.  Note that we don't simply clone the
        # current block.  We need to collect a single set of equations that can be easily 
        # expressed.
        #
        d2 = {}
        B2 = {}
        vtmp = {}
        utmp = {}
        sids_set = set()
        sids_list = []
        #
        block = Block(concrete=True)
        block.u = VarList()
        block.v = VarList()
        block.c1 = ConstraintList()
        block.c2 = ComplementarityList()
        block.c3 = ComplementarityList()
        #
        # Collect submodel objective terms
        #
        for odata in submodel.component_data_objects(Objective, active=True):
            if odata.sense == maximize:
                d_sense = -1
            else:
                d_sense = 1
            #
            # Iterate through the variables in the canonical representation
            #
            o_terms = generate_canonical_repn(odata.expr, compute_values=False)
            for i in range(len(o_terms.variables)):
                var = o_terms.variables[i]
                if var.parent_component().name in self._fixed_upper_vars:
                    #
                    # Skip fixed upper variables
                    #
                    continue
                #
                # Store the coefficient for the variable.  The coefficient is
                # negated if the objective is maximized.
                #
                id_ = id(var)
                d2[id_] = d_sense * o_terms.linear[i]
                if not id_ in sids_set:
                    sids_set.add(id_)
                    sids_list.append(id_)
            # Stop after the first objective
            break
        #
        # Iterate through all lower level variables, adding dual variables 
        # and complementarity slackness conditions for y bound constraints
        #
        for vcomponent in instance.component_objects(Var, active=True):
            if vcomponent.name in self._fixed_upper_vars:
                #
                # Skip fixed upper variables
                #
                continue
            for ndx in vcomponent:
                #
                # For each index, get the bounds for the variable
                #
                lb, ub = vcomponent[ndx].bounds
                if not lb is None:
                    #
                    # Add the complementarity slackness condition for a lower bound
                    #
                    v = block.v.add()
                    block.c3.add( complements(vcomponent[ndx] >= lb, v >= 0) )
                else:
                    v = None
                if not ub is None:
                    #
                    # Add the complementarity slackness condition for an upper bound
                    #
                    w = block.v.add()
                    vtmp[id(vcomponent[ndx])] = w
                    block.c3.add( complements(vcomponent[ndx] <= ub, w >= 0) )
                else:
                    w = None
                if not (v is None and w is None):
                    #
                    # Record the variables for which complementarity slackness conditions
                    # were created.
                    #
                    id_ = id(vcomponent[ndx])
                    vtmp[id_] = (v,w)
                    if not id_ in sids_set:
                        sids_set.add(id_)
                        sids_list.append(id_)
        #
        # Iterate through all constraints, adding dual variables and 
        # complementary slackness conditions (for inequality constraints)
        #
        for cdata in submodel.component_data_objects(Constraint, active=True):
            if cdata.equality:
                # Don't add a complementary slackness condition for an equality constraint
                u = block.u.add()
                utmp[id(cdata)] = (None,u)
            else:
                if not cdata.lower is None:
                    #
                    # Add the complementarity slackness condition for a greater-than inequality
                    #
                    u = block.u.add()
                    block.c2.add( complements(- cdata.body <= - cdata.lower, u >= 0) )
                else:
                    u = None
                if not cdata.upper is None:
                    #
                    # Add the complementarity slackness condition for a less-than inequality
                    #
                    w = block.u.add()
                    block.c2.add( complements(cdata.body <= cdata.upper, w >= 0) )
                else:
                    w = None
                if not (u is None and w is None):
                    utmp[id(cdata)] = (u,w)
            #
            # Store the coefficients for the contraint variables that are not fixed
            #
            c_terms = generate_canonical_repn(cdata.body, compute_values=False)
            for i in range(len(c_terms.variables)):
                var = c_terms.variables[i]
                if var.parent_component().name in self._fixed_upper_vars:
                    continue
                id_ = id(var)
                B2.setdefault(id_,{}).setdefault(id(cdata),c_terms.linear[i])
                if not id_ in sids_set:
                    sids_set.add(id_)
                    sids_list.append(id_)
        #
        # Generate stationarity equations
        #
        tmp__ = (None, None)
        for vid in sids_list:
            exp = d2.get(vid,0)
            #
            lb_dual, ub_dual = vtmp.get(vid, tmp__)
            if vid in vtmp:
                if not lb_dual is None:
                    exp -= lb_dual             # dual for variable lower bound
                if not ub_dual is None:
                    exp += ub_dual             # dual for variable upper bound
            #
            B2_ = B2.get(vid,{})
            utmp_keys = list(utmp.keys())
            if self._deterministic:
                utmp_keys.sort(key=lambda x:utmp[x][0].cname() if utmp[x][1] is None else utmp[x][1].cname())
            for uid in utmp_keys:
                if uid in B2_:
                    lb_dual, ub_dual = utmp[uid]
                    if not lb_dual is None:
                        exp -= B2_[uid] * lb_dual
                    if not ub_dual is None:
                        exp += B2_[uid] * ub_dual
            if type(exp) in six.integer_types or type(exp) is float:
                # TODO: Annotate the model as unbounded
                raise IOError("Unbounded variable without side constraints")
            else:
                block.c1.add( exp == 0 )
        #
        # Return block
        #
        return block

