import pulp

from dipapi import DipAPI

from _dippy import *

from collections import defaultdict

try:
    from gimpy import BBTree
except ImportError:
    gimpy_installed = False
else:
    gimpy_installed = True

etree_installed = False

class DipError(Exception):
    """
    Dip Exception
    """

_Solve = Solve
def Solve(prob, params=None):
    """
    Solve a DipProblem instance, returning a solution object

    @param prob A DipProblem instance to solve
    @param params A dictionary of parameters to pass to DIP
    """

    # params is a dictionary, keys are strings, values are
    # strings or dictionaries

    # if value is a dictionary, key is a section, and items
    # of dictionary are names/values

    # if value is a string, then section is NULL and key is name
    # for these parameters we also assign them to the 'DECOMP'
    # section as a convenience

    # the dictionary is converted into a dictionary of
    # strings indexed by (section, name) tuples

    processed = {}
    if params is None:
        params = {}

    if (prob.branch_method == None) and (prob.display_mode == 'off'):
        params['pyBranchMethod'] = '0'
    if (prob.post_process_node == None) and (prob.display_mode == 'off'):
        params['pyPostProcessNode'] = '0'
    if (prob.post_process_branch == None) and (prob.display_mode == 'off'):
        params['pyPostProcessBranch'] = '0'
    if prob.relaxed_solver == None:
        params['pyRelaxedSolver'] = '0'
    if prob.is_solution_feasible == None:
        params['pyIsSolutionFeasible'] = '0'
    
    if (prob.generate_cuts == None) and (prob.generate_cuts_from_node == None):
        params['pyGenerateCuts'] = '0'

    if prob.generate_cuts != None:
        prob.gen_cuts = True
    else:
        prob.gen_cuts = False
    if prob.generate_cuts_from_node != None:
        prob.gen_cuts_node = True
    else:
        prob.gen_cuts_node = False
    
    if prob.heuristics == None:
        params['pyHeuristics'] = '0'
    if prob.init_vars == None:
        params['pyInitVars'] = '0'
    if prob.is_solution_feasible == None:
        params['pyIsSolutionFeasible'] = '0'

    for key, value in params.items():
        valid_types = (basestring, int, float)
        if not isinstance(key, basestring):
            raise DipError("Bad key in parameter dictionary, expecting string")
        if isinstance(value, dict):
            section = key
            for name, param_val in value.items():
                if not isinstance(param_val, valid_types):
                    raise DipError("Bad value '%s' in parameter dictionary, expecting string or number" % param_val)
                processed[(section, name)] = str(param_val)
        elif isinstance(value, valid_types):
            # add this parameter to both the 'None' section and the 'DECOMP' section
            processed[(None, key)] = str(value)
            processed[('DECOMP', key)] = str(value)
        else:
            raise DipError("Bad value '%s' in parameter dictionary, expecting string" % value)

    # DIP only solves minimisation problems
    if prob.sense == pulp.LpMaximize:
        raise DipError("DIP assumes a minimize objective, but DipProblem has maximize objective.\n" +
                       "Use prob.sense = pulp.LpMinimize and prob.objective *= -1 to remedy")
    
    # DIP only allows non-negative variables. This is difficult
    # to transform automatically, so request re-formulation
    for v in prob.variables():
        if v.lowBound < 0:
            raise DipError("Variable %s has negative lower bound, please re-formulate using sum of non-negative variables" % v.name)
        
    # call the Solve method from _dippy
    try:
        status, message, solList, dualList = _Solve(prob, processed)
        # solList  is a list of (col_name, value) pairs
        # dualList is a list of (row_name, value) pairs

        if prob.Tree.display_mode == 'svg':
          if prob.display_interval is not None:
            prob.Tree.write_as_svg(filename = "%s_%d" % (prob.svg_prefix, prob.last_svg + 1), 
                                   prevfile = "%s_%d" % (prob.svg_prefix, prob.last_svg))
            prob.last_svg += 1
            
    except:
        print "Error returned from _dippy"
        raise

    if solList is None:
        solution = None
    else:
        solDict = dict(solList)
        setVars = set(prob.variables())
        setSolVars = set(solDict.keys())
        diff = setVars.symmetric_difference(setSolVars)
        if len(diff) > 0:
            raise DipError("Solution and variable list don't match in dippy.Solve")

        solution = solDict
        for v in prob.variables():
            v.varValue = solution[v]
        
    if dualList is None:
        duals = None
    else:
        dualDict = dict([(c.getName(), v) for (c, v) in dualList])
        setCons = set(prob.constraints)
        setDualCons = set(dualDict.keys())
        diff = setCons.symmetric_difference(setDualCons)
        if len(diff) > 0:
            raise DipError("Duals and constraint list don't match in dippy.Solve")

        duals = dualDict

    # return status, message, solution and duals
    return status, message, solution, duals

class DipProblem(pulp.LpProblem, DipAPI):

    def __init__(self, *args, **kwargs):
        # callback functions can be passed to class constructor as keyword 
        # arguments
        self.branch_method = kwargs.pop('branch_method', None)
        self.post_process_branch = kwargs.pop('post_process_branch', None)
        self.post_process_node = kwargs.pop('post_process_node', None)
        self.relaxed_solver = kwargs.pop('relaxed_solver', None)
        self.is_solution_feasible = kwargs.pop('is_solution_feasible', None)
        self.generate_cuts = kwargs.pop('generate_cuts', None)
        self.generate_cuts_from_node = kwargs.pop('generate_cuts_from_node', 
                                                  None)
        self.heuristics = kwargs.pop('heuristics', None)
        self.init_vars = kwargs.pop('init_vars', None)
        self.display_mode = kwargs.pop('display_mode', 'off')
        self.display_interval = kwargs.pop('display_interval', 1)
        self.svg_prefix = kwargs.pop('svg_prefix', 'tree')
        
        if self.display_mode != 'off':
            if gimpy_installed == False:
                print "GiMPy not installed. Display mode set to 'off'"
                self.display_mode = 'off'
            else:
                self.Tree = BBTree()
                self.Tree.set_display_mode(self.display_mode)

        super(DipProblem, self).__init__(*args, **kwargs)
        self._subproblem = []
        self.relaxation = RelaxationCollection(self)

    def deepcopy(self):
        # callback functions can be passed to class constructor as keyword 
        # arguments
        dipcopy = DipProblem(name = self.name, sense = self.sense)
        dipcopy.branch_method = self.branch_method
        dipcopy.is_solution_feasible = self.is_solution_feasible
        dipcopy.generate_cuts = self.generate_cuts
        dipcopy.heuristics = self.heuristics
        dipcopy.init_vars = self.init_vars

        # This code is taken from pulp.py and needs to be coordinated
        # with pulp.py to avoid errors
        if dipcopy.objective != None:
            dipcopy.objective = self.objective.copy()
        dipcopy.constraints = {}
        for k,v in self.constraints.iteritems():
            dipcopy.constraints[k] = v.copy()
        dipcopy.sos1 = self.sos1.copy()
        dipcopy.sos2 = self.sos2.copy()

        dipcopy._subproblem = self._subproblem[:]
        for k in self.relaxation.keys():
            dipcopy.relaxation[k] = self.relaxation[k].copy()

        return dipcopy
    
    def variables(self):
        """
        Returns a list of the problem variables
        Overrides LpProblem.variables()
        
        Inputs:
            - none
        
        Returns:
            - A list of the problem variables
        """
        variables = {}
        if self.objective:
            variables.update(self.objective)
        for c in self.constraints.itervalues():
            variables.update(c)
        for p in sorted(self.relaxation.keys()):
            for c in self.relaxation[p].constraints.itervalues():
                variables.update(c)
        variables = list(variables)
        variables = sorted(variables, key=lambda variable: variable.name)
        
        return variables

    def getObjective(self):
        """
        Return objective as a dictionary with LpVariables as keys
        and (non-zero) coefficients as values
        """
        return self.objective

    def getRows(self, problem=None):
        """
        Return constraints as a list of dictionaries with LpVariables as keys
        and (non-zero) coefficients as values. Constraints also have
        getName, getLb and getUb methods (i.e., a LpConstraint)

        problem = None implies the master problem, otherwise problem
        is a subproblem
        """

        if problem is None:
            problem = self

        for n, c in problem.constraints.iteritems():
            if c.name == None:
                c.name = n
        constraints = problem.constraints.values()

        return constraints

    def getCols(self, problem=None):
        """
        Returns a list of variables. Variables have getName, getLb,
        getUb and isInteger methods

        problem = None implies the master problem, otherwise problem
        is a subproblem
        """

        if problem is None:
            variables = self.variables()
        else:
            variables = {}
            for c in problem.constraints.itervalues():
                variables.update(c)
            variables = list(variables)
            variables = sorted(variables, key=lambda variable: variable.name)

        return variables

    def getRelaxsAsDict(self):
        """
        Returns the relaxation subproblems as a dictionary with keys as
        defined by the user and values as subproblems
        """
        return self.relaxation.dict

    def writeRelaxed(self, block, filename, mip = 1):
        """
        Write the given block into a .lp file.
        
        This function writes the specifications (NO objective function,
        constraints, variables) of the defined Lp problem to a file.
        
        Inputs:
            - block -- the key to the block to write
            - filename -- the name of the file to be created.          
                
        Side Effects:
            - The file is created.
        """
        relaxation = self.relaxation[block]
        f = file(filename, "w")
        f.write("\\* "+relaxation.name+" *\\\n")
        f.write("Subject To\n")
        ks = relaxation.constraints.keys()
        ks.sort()
        for k in ks:
            f.write(relaxation.constraints[k].asCplexLpConstraint(k))
        vs = relaxation.variables()
        vs.sort()
        # Bounds on non-"positive" variables
        # Note: XPRESS and CPLEX do not interpret integer variables without 
        # explicit bounds
        if mip:
            vg = [v for v in vs if not (v.isPositive() and \
                                        v.cat == pulp.LpContinuous) \
                and not v.isBinary()]
        else:
            vg = [v for v in vs if not v.isPositive()]
        if vg:
            f.write("Bounds\n")
            for v in vg:
                f.write("%s\n" % v.asCplexLpVariable())
        # Integer non-binary variables
        if mip:
            vg = [v for v in vs if v.cat == pulp.LpInteger and \
                                   not v.isBinary()]
            if vg:
                f.write("Generals\n")
                for v in vg: f.write("%s\n" % v.name)
            # Binary variables
            vg = [v for v in vs if v.isBinary()]
            if vg:
                f.write("Binaries\n")
                for v in vg: f.write("%s\n" % v.name)
        f.write("End\n")
        f.close()
        
    def chooseBranchSet(self, xhat):
        """
        Finds the best branch for a fractional solution

        Inputs:
        xhat (list of (LpVariable, value) tuples) = list of solution values for all variables

        Output:
        down_lb, down_ub, up_lb, up_ub (tuple of (LpVariable, value) dictionaries) =
        lower and upper bounds for down branch, lower and upper bounds for up branch
        """

        if not self.branch_method:
            return None

        xhatDict = dict(xhat)
        setVars = set(self.variables())
        setXhatVars = set(xhatDict.keys())
        diff = setVars.symmetric_difference(setXhatVars)
        if len(diff) > 0:
            raise DipError("Solution and variable list don't match in chooseBranchSet")

        return self.branch_method(self, xhatDict)

    def decipherNode(self, output):
        outputDict = dict(output)
        if "xhat" in outputDict.keys():
            xhat = outputDict["xhat"]
            outputDict["xhat"] = dict(xhat)
        if "bounds" in outputDict.keys():
            bounds = outputDict["bounds"]
            outputDict["bounds"] = dict(bounds)
            
        return outputDict

    def postProcessNode(self, node):
        """
        Returns information from the node that has just been processed.

        Inputs:
        output (list of (parameter, value) tuples) = list of output values 
        from the node
        """
        nodeDict = self.decipherNode(node)
        
        if gimpy_installed:
            nodeInd = nodeDict["nodeIndex"]
            parentInd = nodeDict["parentIndex"]
            nodeQuality = nodeDict["nodeQuality"]
            branchedDir = nodeDict["branchedDir"]
            nodeStatus = nodeDict["nodeStatus"]
            if branchedDir == -1:
                branch_direction = 'R'
            else:
                branch_direction = 'L'

            if nodeStatus == 'Infeasible':
                status = 'I'
                color = 'orange'
            elif nodeStatus == 'Candidate':
                status = 'C'
                color = 'yellow'
            elif nodeStatus == 'Solution':
                status = 'S'
                color = 'lightblue'
            else:
                status = 'P'
                color = 'red'

            if nodeStatus != 'Infeasible':
                label = status + ": " + "%.1f"%nodeQuality
            else:
                label = 'I'

            numNodes = len(self.Tree.get_node_list())
            if parentInd == -1:
                self.Tree.add_root(str(nodeInd), label = label, status = 'C', 
                                   obj = nodeQuality, color = color, 
                                   style = 'filled', fillcolor = color)
                if self.Tree.display_mode == 'svg':
                  if self.display_interval is not None:
                      if numNodes % self.display_interval in [0, 1]:
                        self.Tree.write_as_svg(filename = "%s_0" % self.svg_prefix, 
                                               nextfile = "%s_1" % self.svg_prefix, 
                                               highlight = nodeInd)
                        self.last_svg = 0
                    
                numNodes += 1
            else:
                if branch_direction == 'L':
                    n = self.Tree.get_left_child(str(parentInd))
                else:
                    n = self.Tree.get_right_child(str(parentInd))
                edge_label = self.Tree.get_edge_attr(parentInd, n, 'label')
                self.Tree.del_node(n)
                if branch_direction == 'L':
                    self.Tree.add_left_child(str(nodeInd), str(parentInd), 
                                             label = label, 
                                             status = status, 
                                             obj = nodeQuality, color = color, 
                                             style = 'filled', 
                                             fillcolor = color)
                else:
                    self.Tree.add_right_child(str(nodeInd), str(parentInd), 
                                              label = label, 
                                              status = status, 
                                              obj = nodeQuality, color = color, 
                                              style = 'filled', 
                                              fillcolor = color)
                if edge_label is not None:
                        self.Tree.set_edge_attr(str(parentInd), str(nodeInd), 
                                                'label', edge_label)
                if (self.Tree.display_mode == 'svg'):
                  if self.display_interval is not None:
                      if numNodes % self.display_interval in [0, 1]:
                        self.Tree.write_as_svg(filename = "%s_%d" % (self.svg_prefix, self.last_svg + 1), 
                                               prevfile = "%s_%d" % (self.svg_prefix, self.last_svg), 
                                               nextfile = "%s_%d" % (self.svg_prefix, self.last_svg + 2), 
                                               highlight = nodeInd)
                        self.last_svg += 1
                        
            if self.display_interval is not None:
                if numNodes % self.display_interval in [0, 1]:
                    self.Tree.display()

        if self.post_process_node is not None:
            self.post_process_node(self, nodeDict)

    def postProcessBranch(self, branchInfo):
        """
        Returns information from the node that has just been processed.

        Inputs:
        output (list of (parameter, value) tuples) describing branching decision
        """
        outputDict = dict(branchInfo)
            
        if gimpy_installed:
            nodeInd = outputDict['nodeIndex']
            nodeQuality = outputDict['nodeQuality']
            numNodes = len(self.Tree.get_node_list())
            for n in outputDict:
                if n == 'pDownUB':
                    self.Tree.add_left_child(str(-numNodes), str(nodeInd), 
                                             label = 'C', 
                                             status = 'C', 
                                             obj = nodeQuality, 
                                             color = 'yellow', 
                                             style = 'filled', 
                                             fillcolor = 'yellow')
                    if len(outputDict[n]) == 1:
                        self.Tree.set_edge_attr(str(nodeInd), str(-numNodes), 
                                                'label', 
                                                ''.join(x for x in 
                                                        str(outputDict[n].keys()[0]) 
                                                        if x not in '()_')
                                                + '<=' + 
                                                str(int(outputDict[n].values()[0])))
                    numNodes += 1

                elif n == 'pUpLB':
                    self.Tree.add_right_child(str(-numNodes), str(nodeInd), 
                                              label = 'C', 
                                              status = 'C', 
                                              obj = nodeQuality, 
                                              color = 'yellow', 
                                              style = 'filled', 
                                              fillcolor = 'yellow')
                    if len(outputDict[n]) == 1:
                        self.Tree.set_edge_attr(str(nodeInd), str(-numNodes), 
                                                'label', 
                                                ''.join(x for x in 
                                                        str(outputDict[n].keys()[0]) 
                                                        if x not in '()_')
                                                + '>=' + 
                                                str(int(outputDict[n].values()[0])))
                    numNodes += 1
                    
        self.Tree.set_node_attr(str(nodeInd), 'color', 'green')
        self.Tree.set_node_attr(str(nodeInd), 'fillcolor', 'green')
            
        if self.post_process_branch is not None:
            self.post_process_branch(self, outputDict)

    def solveRelaxed(self, key, redCostX, convexDual):
        """
        Returns solutions to the whichBlock relaxed subproblem

        Inputs: 
        key (Python Object) = key of relaxed subproblem to be solved
        redCostX (list of (LpVariable, value) tuples) = list of reduced costs
        for all variables convexDual (float) = dual for convexity constraint
        for this relaxed subproblem

        Output: 
        varList (list of (cost, reduced cost, (LpVariable, value)
        dictionaries)) = solution for this relaxed subproblem expressed as a
        cost, reduced cost and dictionary of non-zero values for variables
        """

        # transform redCostX into a dictionary
        redCostDict = dict(redCostX)
        setVars = set(self.variables())
        setRedCostVars = set(redCostDict.keys())
        diff = setVars.symmetric_difference(setRedCostVars)
        if len(diff) > 0:
            print diff
            raise DipError("Reduced cost and variable list don't match in",
                           "solveRelaxed")

        return self.relaxed_solver(self, key, redCostDict, convexDual)

    def isUserFeasible(self, sol, tol):
        """
        Lets the user decide if an integer solution is really feasible

        Inputs:
        sol (list of (LpVariable, value) tuples) = list of solution
        values for all variables tol (double) = zero tolerance

        Outputs:
        (boolean) = false if not feasible (generate cuts) or true if feasible
        """
        if not self.is_solution_feasible:
            return None

        solDict = dict(sol)
        setVars = set(self.variables())
        setSolVars = set(solDict.keys())
        diff = setVars.symmetric_difference(setSolVars)
        if len(diff) > 0:
            raise DipError("Solution and variable list don't match in isUserFeasible")

        return self.is_solution_feasible(self, solDict, tol)

    def generateCuts(self, node):
        """
        Lets the user generate cuts to remove fractional "pieces" of xhat

        Inputs:
        node (list of (string, object) tuples) = list of node properties

        Output:
        cutList (list of LpConstraints) =
        cuts for this fractional solution expressed as a list LpConstraints,
        i.e., a dictionary with LpVariables as keys and (non-zero) coefficients
        as values with getName, getLb and getUb bound methods
        """
        if (not self.gen_cuts) and (not self.gen_cuts_node):
            return None

        nodeDict = self.decipherNode(node)
        xhatDict = nodeDict["xhat"]
        setVars = set(self.variables())
        setXhatVars = set(xhatDict.keys())
        diff = setVars.symmetric_difference(setXhatVars)
        if len(diff) > 0:
            raise DipError("Solution and variable list don't match in generateCuts")

        # Generate a list of cuts as LpConstraints
        if self.gen_cuts:
            cuts = self.generate_cuts(self, xhatDict)
        else:
            cuts = []

        if self.gen_cuts_node:
            cuts.extend(self.generate_cuts_from_node(self, nodeDict))
                
        return cuts

    def solveHeuristics(self, xhat, costX):
        """
        Lets the user generate (heuristic) solutions from a fractional solution

        Inputs:
        xhat  (list of (LpVariable, value) tuples) = list of solution values for all variables
        costX (list of (LpVariable, value) tuples) = list of costs for all variables

        Outputs:
        solList (list of (LpVariable, value) dictionaries) =
        solutions found from this fractional solution expressed as a
        dictionary of non-zero values for variables
        """

        if not self.heuristics:
            return None

        # transform xhat into a dictionary
        xhatDict = dict(xhat)
        setVars = set(self.variables())
        setXhatVars = set(xhatDict.keys())
        diff = setVars.symmetric_difference(setXhatVars)
        if len(diff) > 0:
            raise DipError("Solution and variable list don't match in solveHeuristics")

        # transform costs into a dictionary
        costDict = dict(costX)
        setCostVars = set(costDict.keys())
        diff = setVars.symmetric_difference(setCostVars)
        if len(diff) > 0:
            raise DipError("Cost and variable list don't match in solveHeuristics")

        return self.heuristics(self, xhatDict, costDict)

    def generateInitVars(self):
        """
        Returns initial solutions to relaxed subproblems

        Inputs:
        None

        Output:
        varList (list of (subproblem key, (cost, (LpVariable, value) dictionaries))) =
        initial solutions for the relaxed subproblems expressed as a cost and
        dictionary of non-zero values for variables
        """

        if not self.init_vars:
            return None

        return self.init_vars(self)
    
class RelaxationCollection(object):
    """
    A simple defaultdict for holding relaxation problems
    """
    PROBLEM_CLASS = pulp.LpProblem


    def __init__(self, parent):
        self.parent = parent
        self.dict = {}

    def __getitem__(self, name):
        if name not in self.dict:
            self.dict[name] = self.PROBLEM_CLASS()
        return self.dict[name]

    def __setitem__(self, name, value):
        self.dict[name] = value

    def keys(self):
        return self.dict.keys()

    def values(self):
        return self.dict.values()



