text: from __future__ import division
  import copy
  import itertools as it
  
  import numpy as np
  import networkx as nx
  
  from pgmpy.extern.six.moves import filter, range
  from pgmpy.inference import Inference
  from pgmpy.models import MarkovModel
  from pgmpy.factors import Factor
  
  
  class Mplp(Inference):
      """
      Class for performing approximate inference using Max-Product Linear Programming method.
  
      We derive message passing updates that result in monotone decrease of the dual of the
      MAP LP Relaxation.
  
      Parameters
      ----------
      model: MarkovModel for which inference is to be performed.
      Examples
      --------
      >>> from pgmpy.models import MarkovModel
      >>> from pgmpy.factors import Factor
      >>> import numpy as np
      >>> student = MarkovModel()
      >>> student.add_edges_from([('A', 'B'), ('B', 'C'), ('C', 'D'), ('E', 'F')])
      >>> factor_a = Factor(['A'], cardinality=[2], value=np.array([0.54577, 1.8323]))
      >>> factor_b = Factor(['B'], cardinality=[2], value=np.array([0.93894, 1.065]))
      >>> factor_c = Factor(['C'], cardinality=[2], value=np.array([0.89205, 1.121]))
      >>> factor_d = Factor(['D'], cardinality=[2], value=np.array([0.56292, 1.7765]))
      >>> factor_e = Factor(['E'], cardinality=[2], value=np.array([0.47117, 2.1224]))
      >>> factor_f = Factor(['F'], cardinality=[2], value=np.array([1.5093, 0.66257]))
      >>> factor_a_b = Factor(['A', 'B'], cardinality=[2, 2], value=np.array([1.3207, 0.75717, 0.75717, 1.3207]))
      >>> factor_b_c = Factor(['B', 'C'], cardinality=[2, 2], value=np.array([0.00024189, 4134.2, 4134.2, 0.00024189]))
      >>> factor_c_d = Factor(['C', 'D'], cardinality=[2, 2], value=np.array([0.0043227, 231.34, 231.34, 0.0043227]))
      >>> factor_d_e = Factor(['E', 'F'], cardinality=[2, 2], value=np.array([31.228, 0.032023, 0.032023, 31.228]))
      >>> student.add_factors(factor_a, factor_b, factor_c, factor_d, factor_e, factor_f, factor_a_b,
      ...    factor_b_c, factor_c_d, factor_d_e)
      >>> mplp = Mplp(student)
      """
      def __init__(self, model):
          if not isinstance(model, MarkovModel):
              raise TypeError('Only MarkovModel is supported')
  
          super(Mplp, self).__init__(model)
          self.model = model
  
          # S = \{c \cap c^{'} : c, c^{'} \in C, c \cap c^{'} \neq \emptyset\}
          self.intersection_set_variables = set()
          # We generate the Intersections of all the pairwise edges taken one at a time to form S
          for edge_pair in it.combinations(model.edges(), 2):
              self.intersection_set_variables.add(frozenset(edge_pair[0]) & frozenset(edge_pair[1]))
  
          # The corresponding optimization problem = \min_{\delta}{dual_lp(\delta)} where:
          # dual_lp(\delta) = \sum_{i \in V}{max_{x_i}(Objective[nodes])} + \sum_{f /in F}{max_{x_f}(Objective[factors])
          # Objective[nodes] = \theta_i(x_i) + \sum_{f \mid i \in f}{\delta_{fi}(x_i)}
          # Objective[factors] = \theta_f(x_f) - \sum_{i \in f}{\delta_{fi}(x_i)}
          # In a way Objective stores the corresponding optimization problem for all the nodes and the factors.
  
          # Form Objective and cluster_set in the form of a dictionary.
          self.objective = {}
          self.cluster_set = {}
          for factor in model.get_factors():
              scope = frozenset(factor.scope())
              self.objective[scope] = factor
              # For every factor consisting of more that a single node, we initialize a cluster.
              if len(scope) > 1:
                  self.cluster_set[scope] = self.Cluster(self.intersection_set_variables, factor)
  
          # dual_lp(\delta) is the dual linear program
          self.dual_lp = sum([np.amax(self.objective[obj].values) for obj in self.objective])
  
          # Best integral value of the primal objective is stored here
          self.best_int_objective = 0
  
          # Assignment of the nodes that results in the "maximum" integral value of the primal objective
          self.best_assignment = {}
          # Results of the "maximum" integral value of the primal objective.
          self.best_decoded_result = {}
          # This sets the minimum width between the dual objective decrements. Default value = 0.0002. This can be
          # changed in the map_query() method.
          self.dual_threshold = 0.0002
          # This sets the threshold for the integrality gap below which we say that the solution is satisfactory.
          # Default value = 0.0002. This can be changed in the map_query() method.
          self.integrality_gap_threshold = 0.0002
  
      class Cluster(object):
          """
          Inner class for representing a cluster.
          A cluster is a subset of variables.
  
          Parameters
          ----------
          set_of_variables: tuple
                            This is the set of variables that form the cluster.
  
          intersection_set_variables: set containing frozensets.
                                      collection of intersection of all pairs of cluster variables.
                          For eg: \{\{C_1 \cap C_2\}, \{C_2 \cap C_3\}, \{C_3 \cap C_1\} \} for clusters C_1, C_2 & C_3.
  
          cluster_potential: Factor
                             Each cluster has a initial probability distribution provided beforehand.
          """
          def __init__(self, intersection_set_variables, cluster_potential):
              """
              Initialization of the current cluster
              """
  
              # The variables with which the cluster is made of.
              self.cluster_variables = frozenset(cluster_potential.scope())
  
              # The cluster potentials must be specified before only.
              self.cluster_potential = copy.deepcopy(cluster_potential)
  
              # Generate intersection sets for this cluster; S(c)
              self.intersection_sets_for_cluster_c = [intersect.intersection(self.cluster_variables)
                                                      for intersect in intersection_set_variables
                                                      if intersect.intersection(self.cluster_variables)]
  
              # Initialize messages from this cluster to its respective intersection sets
              # \lambda_{c \rightarrow \s} = 0
              self.message_from_cluster = {}
              for intersection in self.intersection_sets_for_cluster_c:
                  # Present variable. It can be a node or an edge too. (that is ['A'] or ['A', 'C'] too)
                  present_variables = list(intersection)
  
                  # Present variables cardinality
                  present_variables_card = cluster_potential.get_cardinality(present_variables)
                  present_variables_card = [present_variables_card[var] for var in present_variables]
  
                  # We need to create a new factor whose messages are blank
                  self.message_from_cluster[intersection] = \
                      Factor(present_variables, present_variables_card, np.zeros(np.prod(present_variables_card)))
  
      def _update_message(self, sending_cluster):
  
          """
          This is the message-update method.
  
          Parameters
          ----------
          sending_cluster: The resulting messages are lambda_{c-->s} from the given
              cluster 'c' to all of its intersection_sets 's'.
              Here 's' are the elements of intersection_sets_for_cluster_c.
  
          Reference
          ---------
          Fixing Max-Product: Convergent Message-Passing Algorithms for MAP LP Relaxations
          by Amir Globerson and Tommi Jaakkola.
          Section 6, Page: 5; Beyond pairwise potentials: Generalized MPLP
          Later Modified by Sontag in "Introduction to Dual decomposition for Inference" Pg: 7 & 17
          """
  
          # The new updates will take place for the intersection_sets of this cluster.
          # The new updates are:
          # \delta_{f \rightarrow i}(x_i) = - \delta_i^{-f} +
          # 1/{\| f \|} max_{x_{f-i}}\left[{\theta_f(x_f) + \sum_{i' in f}{\delta_{i'}^{-f}}(x_i')} \right ]
  
          # Step. 1) Calculate {\theta_f(x_f) + \sum_{i' in f}{\delta_{i'}^{-f}}(x_i')}
          objective_cluster = self.objective[sending_cluster.cluster_variables]
          for current_intersect in sending_cluster.intersection_sets_for_cluster_c:
              objective_cluster += self.objective[current_intersect]
  
          updated_results = []
          objective = []
          for current_intersect in sending_cluster.intersection_sets_for_cluster_c:
              # Step. 2) Maximize step.1 result wrt variables present in the cluster but not in the current intersect.
              phi = objective_cluster.maximize(list(sending_cluster.cluster_variables - current_intersect),
                                               inplace=False)
  
              # Step. 3) Multiply 1/{\| f \|}
              intersection_length = len(sending_cluster.intersection_sets_for_cluster_c)
              phi *= (1 / intersection_length)
              objective.append(phi)
  
              # Step. 4) Subtract \delta_i^{-f}
              # These are the messages not emanating from the sending cluster but going into the current intersect.
              # which is = Objective[current_intersect_node] - messages from the cluster to the current intersect node.
              updated_results.append(phi + -1 * (self.objective[current_intersect] + -1 * sending_cluster.
                                                 message_from_cluster[current_intersect]))
  
          # This loop is primarily for simultaneous updating:
          # 1. This cluster's message to each of the intersects.
          # 2. The value of the Objective for intersection_nodes.
          index = -1
          cluster_potential = copy.deepcopy(sending_cluster.cluster_potential)
          for current_intersect in sending_cluster.intersection_sets_for_cluster_c:
              index += 1
              sending_cluster.message_from_cluster[current_intersect] = updated_results[index]
              self.objective[current_intersect] = objective[index]
              cluster_potential += (-1) * updated_results[index]
  
          # Here we update the Objective for the current factor.
          self.objective[sending_cluster.cluster_variables] = cluster_potential
  
      def _local_decode(self):
          """
          Finds the index of the maximum values for all the single node dual objectives.
  
          Reference:
          code presented by Sontag in 2012 here: http://cs.nyu.edu/~dsontag/code/README_v2.html
          """
          # The current assignment of the single node factors is stored in the form of a dictionary
          decoded_result_assignment = {node: np.argmax(self.objective[node].values)
                                       for node in self.objective if len(node) == 1}
          # Use the original cluster_potentials of each factor to find the primal integral value.
          # 1. For single node factors
          integer_value = sum([self.factors[variable][0].values[decoded_result_assignment[frozenset([variable])]]
                               for variable in self.variables])
          # 2. For clusters
          for cluster_key in self.cluster_set:
              cluster = self.cluster_set[cluster_key]
              index = [tuple([variable, decoded_result_assignment[frozenset([variable])]])
                       for variable in cluster.cluster_variables]
              integer_value += cluster.cluster_potential.reduce(index, inplace=False).values
  
          # Check if this is the best assignment till now
          if self.best_int_objective < integer_value:
              self.best_int_objective = integer_value
              self.best_assignment = decoded_result_assignment
  
      def _is_converged(self, dual_threshold=None, integrality_gap_threshold=None):
          """
          This method checks the integrality gap to ensure either:
              * we have found a near to exact solution or
              * stuck on a local minima.
  
          Parameters
          ----------
          dual_threshold: double
                          This sets the minimum width between the dual objective decrements. If the decrement is lesser
                          than the threshold, then that means we have stuck on a local minima.
  
          integrality_gap_threshold: double
                                     This sets the threshold for the integrality gap below which we say that the solution
                                     is satisfactory.
  
          References
          ----------
          code presented by Sontag in 2012 here: http://cs.nyu.edu/~dsontag/code/README_v2.html
          """
          # Find the new objective after the message updates
          new_dual_lp = sum([np.amax(self.objective[obj].values) for obj in self.objective])
  
          # Update the dual_gap as the difference between the dual objective of the previous and the current iteration.
          self.dual_gap = abs(self.dual_lp - new_dual_lp)
  
          # Update the integrality_gap as the difference between our best result vs the dual objective of the lp.
          self.integrality_gap = abs(self.dual_lp - self.best_int_objective)
  
          # As the decrement of the dual_lp gets very low, we assume that we might have stuck in a local minima.
          if dual_threshold and self.dual_gap < dual_threshold:
              return True
          # Check the threshold for the integrality gap
          elif integrality_gap_threshold and self.integrality_gap < integrality_gap_threshold:
              return True
          else:
              self.dual_lp = new_dual_lp
              return False
  
      def find_triangles(self):
          """
          Finds all the triangles present in the given model
  
          Examples
          --------
          >>> from pgmpy.models import MarkovModel
          >>> from pgmpy.factors import Factor
          >>> from pgmpy.inference import Mplp
          >>> mm = MarkovModel()
          >>> mm.add_nodes_from(['x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7'])
          >>> mm.add_edges_from([('x1', 'x3'), ('x1', 'x4'), ('x2', 'x4'),
          ...                    ('x2', 'x5'), ('x3', 'x6'), ('x4', 'x6'),
          ...                    ('x4', 'x7'), ('x5', 'x7')])
          >>> phi = [Factor(edge, [2, 2], np.random.rand(4)) for edge in mm.edges()]
          >>> mm.add_factors(*phi)
          >>> mplp = Mplp(mm)
          >>> mplp.find_triangles()
          """
          return list(filter(lambda x: len(x) == 3, nx.find_cliques(self.model)))
  
      def _update_triangles(self, triangles_list):
          """
          From a set of variables forming a triangle in the model, we form the corresponding Clusters.
          These clusters are then appended to the code.
  
          Parameters
          ----------
          triangle_list : list
                          The list of variables forming the triangles to be updated. It is of the form of
                          [['var_5', 'var_8', 'var_7'], ['var_4', 'var_5', 'var_7']]
  
          """
          new_intersection_set = []
          for triangle_vars in triangles_list:
              cardinalities = [self.cardinality[variable] for variable in triangle_vars]
              current_intersection_set = [frozenset(intersect) for intersect in it.combinations(triangle_vars, 2)]
              current_factor = Factor(triangle_vars, cardinalities, np.zeros(np.prod(cardinalities)))
              self.cluster_set[frozenset(triangle_vars)] = self.Cluster(current_intersection_set, current_factor)
              # add new factors
              self.model.factors.append(current_factor)
              # add new intersection sets
              new_intersection_set.extend(current_intersection_set)
              # add new factors in objective
              self.objective[frozenset(triangle_vars)] = current_factor
  
      def _get_triplet_scores(self, triangles_list):
          """
          Returns the score of each of the triplets found in the current model
  
          Parameters
          ---------
          triangles_list: list
                          The list of variables forming the triangles to be updated. It is of the form of
                          [['var_5', 'var_8', 'var_7'], ['var_4', 'var_5', 'var_7']]
  
          Return: {frozenset({'var_8', 'var_5', 'var_7'}): 5.024, frozenset({'var_5', 'var_4', 'var_7'}): 10.23}
          """
          triplet_scores = {}
          for triplet in triangles_list:
  
              # Find the intersection sets of the current triplet
              triplet_intersections = [intersect for intersect in it.combinations(triplet, 2)]
  
              # Independent maximization
              ind_max = sum([np.amax(self.objective[frozenset(intersect)].values) for intersect in triplet_intersections])
  
              # Joint maximization
              joint_max = self.objective[frozenset(triplet_intersections[0])]
              for intersect in triplet_intersections[1:]:
                  joint_max += self.objective[frozenset(intersect)]
              joint_max = np.amax(joint_max.values)
              # score = Independent maximization solution - Joint maximization solution
              score = ind_max - joint_max
              triplet_scores[frozenset(triplet)] = score
  
          return triplet_scores
  
      def _run_mplp(self, no_iterations):
          """
          Updates messages until either Mplp converges or if doesn't converges; halts after no_iterations.
  
          Parameters
          --------
          no_iterations:  integer
                          Number of maximum iterations that we want MPLP to run.
          """
          for niter in range(no_iterations):
              # We take the clusters in the order they were added in the model and update messages for all factors whose
              # scope is greater than 1
              for factor in self.model.get_factors():
                  if len(factor.scope()) > 1:
                      self._update_message(self.cluster_set[frozenset(factor.scope())])
              # Find an integral solution by locally maximizing the single node beliefs
              self._local_decode()
              # If mplp converges to a global/local optima, we break.
              if self._is_converged(self.dual_threshold, self.integrality_gap_threshold) and niter >= 16:
                  break
  
      def _tighten_triplet(self, max_iterations, later_iter, max_triplets, prolong):
          """
          This method finds all the triplets that are eligible and adds them iteratively in the bunch of max_triplets
  
          Parameters
          ----------
          max_iterations: integer
                          Maximum number of times we tighten the relaxation
  
          later_iter: integer
                      Number of maximum iterations that we want MPLP to run. This is lesser than the initial number
                      of iterations.
  
          max_triplets: integer
                        Maximum number of triplets that can be added atmost in one iteration.
  
          prolong: bool
                  It sets the continuation of tightening after all the triplets are exhausted
          """
          # Find all the triplets that are possible in the present model
          triangles = self.find_triangles()
          # Evaluate scores for each of the triplets found above
          triplet_scores = self._get_triplet_scores(triangles)
          # Arrange the keys on the basis of increasing order of the values of the dict. triplet_scores
          sorted_scores = sorted(triplet_scores, key=triplet_scores.get)
          for niter in range(max_iterations):
              if self._is_converged(integrality_gap_threshold=self.integrality_gap_threshold):
                  break
              # add triplets that are yet not added.
              add_triplets = []
              for triplet_number in (range(len(sorted_scores))):
                  # At once, we can add atmost 5 triplets
                  if triplet_number >= max_triplets:
                      break
                  add_triplets.append(sorted_scores.pop())
              # Break from the tighten triplets loop if there are no triplets to add if the prolong is set to False
              if not add_triplets and prolong is False:
                      break
              # Update the eligible triplets to tighten the relaxation
              self._update_triangles(add_triplets)
              # Run MPLP for a maximum of later_iter times.
              self._run_mplp(later_iter)
  
      def get_integrality_gap(self):
          """
          Returns the integrality gap of the current state of the Mplp algorithm. The lesser it is, the closer we are
                  towards the exact solution.
  
          Example:
          --------
          >>> from pgmpy.models import MarkovModel
          >>> from pgmpy.factors import Factor
          >>> from pgmpy.inference import Mplp
          >>> mm = MarkovModel()
          >>> mm.add_nodes_from(['x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7'])
          >>> mm.add_edges_from([('x1', 'x3'), ('x1', 'x4'), ('x2', 'x4'),
          ...                    ('x2', 'x5'), ('x3', 'x6'), ('x4', 'x6'),
          ...                    ('x4', 'x7'), ('x5', 'x7')])
          >>> phi = [Factor(edge, [2, 2], np.random.rand(4)) for edge in mm.edges()]
          >>> mm.add_factors(*phi)
          >>> mplp = Mplp(mm)
          >>> mplp.map_query()
          >>> int_gap = mplp.get_integrality_gap()
          """
  
          return self.integrality_gap
  
      def query(self):
          raise NotImplementedError("map_query() is the only query method available.")
  
      def map_query(self, init_iter=1000, later_iter=20, dual_threshold=0.0002, integrality_gap_threshold=0.0002,
                    tighten_triplet=True, max_triplets=5, max_iterations=100, prolong=False):
          """
          MAP query method using Max Product LP method.
          This returns the best assignment of the nodes in the form of a dictionary.
  
          Parameters
          ----------
          init_iter: integer
                     Number of maximum iterations that we want MPLP to run for the first time.
  
          later_iter: integer
                      Number of maximum iterations that we want MPLP to run for later iterations
  
          dual_threshold: double
                          This sets the minimum width between the dual objective decrements. If the decrement is lesser
                          than the threshold, then that means we have stuck on a local minima.
  
          integrality_gap_threshold: double
                                     This sets the threshold for the integrality gap below which we say that the solution
                                     is satisfactory.
  
          tighten_triplet: bool
                           set whether to use triplets as clusters or not.
  
          max_triplets: integer
                        Set the maximum number of triplets that can be added at once.
  
          max_iterations: integer
                          Maximum number of times we tighten the relaxation. Used only when tighten_triplet is set True.
  
          prolong: bool
                   If set False: The moment we exhaust of all the triplets the tightening stops.
                   If set True: The tightening will be performed max_iterations number of times irrespective of the
                                triplets.
  
          Reference:
          Section 3.3: The Dual Algorithm; Tightening LP Relaxation for MAP using Message Passing (2008)
          By Sontag Et al.
          Examples
          --------
          >>> from pgmpy.models import MarkovModel
          >>> from pgmpy.factors import Factor
          >>> import numpy as np
          >>> student = MarkovModel()
          >>> student.add_edges_from([('A', 'B'), ('B', 'C'), ('C', 'D'), ('E', 'F')])
          >>> factor_a = Factor(['A'], cardinality=[2], value=np.array([0.54577, 1.8323]))
          >>> factor_b = Factor(['B'], cardinality=[2], value=np.array([0.93894, 1.065]))
          >>> factor_c = Factor(['C'], cardinality=[2], value=np.array([0.89205, 1.121]))
          >>> factor_d = Factor(['D'], cardinality=[2], value=np.array([0.56292, 1.7765]))
          >>> factor_e = Factor(['E'], cardinality=[2], value=np.array([0.47117, 2.1224]))
          >>> factor_f = Factor(['F'], cardinality=[2], value=np.array([1.5093, 0.66257]))
          >>> factor_a_b = Factor(['A', 'B'], cardinality=[2, 2], value=np.array([1.3207, 0.75717, 0.75717, 1.3207]))
          >>> factor_b_c = Factor(['B', 'C'], cardinality=[2, 2], value=np.array([0.00024189, 4134.2, 4134.2, 0.0002418]))
          >>> factor_c_d = Factor(['C', 'D'], cardinality=[2, 2], value=np.array([0.0043227, 231.34, 231.34, 0.0043227]))
          >>> factor_d_e = Factor(['E', 'F'], cardinality=[2, 2], value=np.array([31.228, 0.032023, 0.032023, 31.228]))
          >>> student.add_factors(factor_a, factor_b, factor_c, factor_d, factor_e, factor_f, factor_a_b,
          ...    factor_b_c, factor_c_d, factor_d_e)
          >>> mplp = Mplp(student)
          >>> result = mplp.map_query()
          Return: {'B': 0.93894, 'C': 1.121, 'A': 1.8323, 'F': 1.5093, 'D': 1.7765, 'E': 2.12239}
          """
  
          self.dual_threshold = dual_threshold
          self.integrality_gap_threshold = integrality_gap_threshold
          # Run MPLP initially for a maximum of init_iter times.
          self._run_mplp(init_iter)
          # If triplets are to be used for the tightening, we proceed as follows
          if tighten_triplet:
              self._tighten_triplet(max_iterations, later_iter, max_triplets, prolong)
          # Get the best result from the best assignment
          self.best_decoded_result = {factor.scope()[0]: factor.values[self.best_assignment[frozenset(factor.scope())]]
                                      for factor in self.model.factors if len(factor.scope()) == 1}
          return self.best_decoded_result
  