text: # Copyright (c) Microsoft Corporation and contributors.
  # Licensed under the MIT License.
  
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  from joblib import Parallel, delayed
  from sklearn.base import BaseEstimator
  from sklearn.utils import check_array, check_random_state, column_or_1d
  from typing_extensions import Literal
  
  from .qap import quadratic_assignment
  
  # Type aliases
  PaddingType = Literal["adopted", "naive"]
  InitMethodType = Literal["barycenter", "rand", "randomized"]
  RandomStateType = Optional[Union[int, np.random.RandomState, np.random.Generator]]
  ArrayLikeOfIndexes = Union[List[int], np.ndarray]
  
  
  class GraphMatch(BaseEstimator):
      """
      This class solves the Graph Matching Problem and the Quadratic Assignment Problem
      (QAP) through an implementation of the Fast Approximate QAP Algorithm (FAQ) (these
      two problems are the same up to a sign change) [1].
  
      This algorithm can be thought of as finding an alignment of the vertices of two
      graphs which minimizes the number of induced edge disagreements, or, in the case
      of weighted graphs, the sum of squared differences of edge weight disagreements.
      The option to add seeds (known vertex correspondence between some nodes) is also
      available [2].
  
  
      Parameters
      ----------
  
      n_init : int, positive (default = 1)
          Number of random initializations of the starting permutation matrix that
          the FAQ algorithm will undergo.
  
      init : string (default = 'barycenter') or 2d-array
          The initial position chosen.
  
          "barycenter" : the non-informative “flat doubly stochastic matrix,”
          :math:`J=1 \\times 1^T /n` , i.e the barycenter of the feasible region. This can
          be thought of as the doubly stochastic matrix from which any permutation is
          equally likely.
  
          "rand" : some random point near :math:`J, (J+K)/2`, where K is some random
          doubly stochastic matrix.
  
          If 2d-array, ``init`` must be :math:`m' x m'`, where :math:`m'` is the number of
          nonseeded vertices of ``B``. This initial position repesents a permutation or
          assignment of the nonseeded vertices of ``B``, and thus must be doubly
          stochastic (all of its rows and columns must sum to 1). Note that if using
          seeds, this permutation/assignment is taken with respect to the nonseeded
          vertices in the order in which they were input, even when
          ``shuffle_input = True``.
  
      max_iter : int, positive (default = 30)
          Integer specifying the max number of Franke-Wolfe iterations.
          FAQ typically converges with modest number of iterations.
  
      shuffle_input : bool (default = True)
          Gives users the option to shuffle the nodes of A matrix to avoid results
          from inputs that were already matched.
  
      eps : float (default = 0.03)
          A positive, threshold stopping criteria such that FW continues to iterate
          while Frobenius norm of :math:`(P_{i}-P_{i+1}) > \\text{eps}`
  
      gmp : bool (default = True)
          Gives users the option to solve QAP rather than the Graph Matching Problem
          (GMP). This is accomplished through a negation of the objective function.
  
      padding : string (default = 'adopted')
          Allows user to specify padding scheme if `A` and `B` are not of equal size.
          Say that `A` and `B` have :math:`n_1` and :math:`n_2` nodes, respectively, and
          :math:`n_1 < n_2`.
  
          "adopted" : matches `A` to the best fitting induced subgraph of `B`. Reduces the
          affinity between isolated vertices added to `A` through padding and low-density
          subgraphs of `B`.
  
          "naive" : matches `A` to the best fitting subgraph of `B`.
  
      random_state : {None, int, `~np.random.RandomState`, `~np.random.Generator`}
          This parameter defines the object to use for drawing random
          variates.
          If `random_state` is ``None`` the `~np.random.RandomState` singleton is
          used.
          If `random_state` is an int, a new ``RandomState`` instance is used,
          seeded with `random_state`.
          If `random_state` is already a ``RandomState`` or ``Generator``
          instance, then that object is used.
          Default is None.
  
      n_jobs : int or None, (default = None)
          The number of jobs to run in parallel. Parallelization is over the
          initializations, so only relevant when ``n_init > 1``. None means 1 unless in a
          joblib.parallel_backend context. -1 means using all processors. See
          :class:`joblib.Parallel` for more details.
  
      Attributes
      ----------
  
      perm_inds_ : array, size (n,) where n is the number of vertices in the fitted graphs.
          The indices of the optimal permutation (with the fixed seeds given) on the nodes of B,
          to best minimize the objective function :math:`f(P) = \\text{trace}(A^T PBP^T)
          + \\text{trace}(S^T P)`.
  
  
      score_ : float
          The objective function value of for the optimal permutation found.
  
      n_iter_ : int
          Number of Frank-Wolfe iterations run. If ``n_init > 1``, :attr:`n_iter_` reflects the number of
          iterations performed at the initialization returned.
  
  
      References
      ----------
      .. [1] J.T. Vogelstein, J.M. Conroy, V. Lyzinski, L.J. Podrazik, S.G. Kratzer,
          E.T. Harley, D.E. Fishkind, R.J. Vogelstein, and C.E. Priebe, “Fast
          approximate quadratic programming for graph matching,” PLOS one, vol. 10,
          no. 4, p. e0121002, 2015.
  
      .. [2] D. Fishkind, S. Adali, H. Patsolic, L. Meng, D. Singh, V. Lyzinski, C. Priebe,
          Seeded graph matching, Pattern Recognit. 87 (2019) 203–215
  
  
  
      """
  
      init: Union[InitMethodType, np.ndarray]
  
      def __init__(
          self,
          n_init: int = 1,
          init: Union[InitMethodType, np.ndarray] = "barycenter",
          max_iter: int = 30,
          shuffle_input: bool = True,
          eps: float = 0.03,
          gmp: bool = True,
          padding: PaddingType = "adopted",
          random_state: RandomStateType = None,
          n_jobs: Optional[int] = None,
      ):
          if type(n_init) is int and n_init > 0:
              self.n_init = n_init
          else:
              msg = '"n_init" must be a positive integer'
              raise TypeError(msg)
          if isinstance(init, np.ndarray):
              self.init = init
          elif init == "rand":
              self.init = "randomized"
          elif init == "barycenter":
              self.init = "barycenter"
          else:
              msg = 'Invalid "init" parameter string'
              raise ValueError(msg)
          if max_iter > 0 and type(max_iter) is int:
              self.max_iter = max_iter
          else:
              msg = '"max_iter" must be a positive integer'
              raise TypeError(msg)
          if type(shuffle_input) is bool:
              self.shuffle_input = shuffle_input
          else:
              msg = '"shuffle_input" must be a boolean'
              raise TypeError(msg)
          if eps > 0 and type(eps) is float:
              self.eps = eps
          else:
              msg = '"eps" must be a positive float'
              raise TypeError(msg)
          if type(gmp) is bool:
              self.gmp = gmp
          else:
              msg = '"gmp" must be a boolean'
              raise TypeError(msg)
          if isinstance(padding, str) and padding in {"adopted", "naive"}:
              self.padding = padding
          elif isinstance(padding, str):
              msg = 'Invalid "padding" parameter string'
              raise ValueError(msg)
          else:
              msg = '"padding" parameter must be of type string'
              raise TypeError(msg)
          self.random_state = random_state
          self.n_jobs = n_jobs
  
      def fit(
          self,
          A: np.ndarray,
          B: np.ndarray,
          seeds_A: ArrayLikeOfIndexes = [],
          seeds_B: ArrayLikeOfIndexes = [],
          S: Optional[np.ndarray] = None,
      ) -> "GraphMatch":
          """
          Fits the model with two assigned adjacency matrices
  
          Parameters
          ----------
          A : 2d-array, square
              A square adjacency matrix
  
          B : 2d-array, square
              A square adjacency matrix
  
          seeds_A : 1d-array, shape (m , 1) where m <= number of nodes (default = [])
              An array where each entry is an index of a node in ``A``.
  
          seeds_B : 1d-array, shape (m , 1) where m <= number of nodes (default = [])
              An array where each entry is an index of a node in `B` The elements of
              ``seeds_A`` and ``seeds_B`` are vertices which are known to be matched, that is,
              ``seeds_A[i]`` is matched to vertex ``seeds_B[i]``.
          S : 2d-array, square
              A similarity matrix. Should be shape (n_A , n_B), the number of nodes in
              ``A`` and ``B``, respectively.
  
              Note: the scale of ``S`` affects the weight placed on the term
              :math:`\\text{trace}(S^T P)` relative to :math:`\\text{trace}(A^T PBP^T)`
              during the optimization process.
  
          Returns
          -------
          self : returns an instance of self
          """
          A = check_array(A, copy=True, ensure_2d=True)
          B = check_array(B, copy=True, ensure_2d=True)
          seeds_A = column_or_1d(seeds_A)
          seeds_B = column_or_1d(seeds_B)
          partial_match = np.column_stack((seeds_A, seeds_B))
  
          if S is None:
              S = np.zeros((A.shape[0], B.shape[1]))
          S = np.atleast_2d(S)
  
          msg = None
          if S.ndim != 2:
              msg = "`S` must have exactly two dimensions"
          elif A.shape[0] != S.shape[0] or B.shape[0] != S.shape[1]:
              msg = "`S` must be of shape (n_A, n_B)"
          if msg is not None:
              raise ValueError(msg)
  
          # pads A and B according to section 2.5 of [2]
          if A.shape[0] != B.shape[0]:
              A, B, S = _adj_pad(A, B, S, self.padding)
  
          options = {
              "maximize": self.gmp,
              "partial_match": partial_match,
              "S": S,
              "P0": self.init,
              "shuffle_input": self.shuffle_input,
              "maxiter": self.max_iter,
              "tol": self.eps,
          }
  
          rng = check_random_state(self.random_state)
          results = Parallel(n_jobs=self.n_jobs)(
              delayed(quadratic_assignment)(A, B, options={**options, **{"rng": r}})
              for r in rng.randint(np.iinfo(np.int32).max, size=self.n_init)
          )
          func = max if self.gmp else min
          res = func(
              results,
              key=lambda x: x.fun,
          )
  
          self.perm_inds_ = res.col_ind  # permutation indices
          self.score_ = res.fun  # objective function value
          self.n_iter_ = res.nit
          return self
  
      def fit_predict(
          self,
          A: np.ndarray,
          B: np.ndarray,
          seeds_A: ArrayLikeOfIndexes = [],
          seeds_B: ArrayLikeOfIndexes = [],
          S: Optional[np.ndarray] = None,
      ) -> np.ndarray:
          """
          Fits the model with two assigned adjacency matrices, returning optimal
          permutation indices
  
          Parameters
          ----------
          A : 2d-array, square
              A square adjacency matrix
  
          B : 2d-array, square
              A square adjacency matrix
  
          seeds_A : 1d-array, shape (m , 1) where m <= number of nodes (default = [])
              An array where each entry is an index of a node in ``A``.
  
          seeds_B : 1d-array, shape (m , 1) where m <= number of nodes (default = [])
              An array where each entry is an index of a node in ``B`` The elements of
              ``seeds_A`` and ``seeds_B`` are vertices which are known to be matched, that is,
              ``seeds_A[i]`` is matched to vertex ``seeds_B[i]``.
  
          S : 2d-array, square
              A similarity matrix. Should be shape (n_A , n_B), the number of nodes in
              ``A`` and ``B``, respectively.
  
              Note: the scale of `S` affects the weight placed on the term
              :math:`\\text{trace}(S^T P)` relative to :math:`\\text{trace}(A^T PBP^T)`
              during the optimization process.
  
          Returns
          -------
          perm_inds_ : 1-d array, some shuffling of [0, n_vert)
              The optimal permutation indices to minimize the objective function
          """
          self.fit(A, B, seeds_A, seeds_B, S)
          return self.perm_inds_
  
  
  def _adj_pad(
      A: np.ndarray, B: np.ndarray, S: np.ndarray, method: PaddingType
  ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
      def pad(X: np.ndarray, n: np.ndarray) -> np.ndarray:
          X_pad = np.zeros((n[1], n[1]))
          X_pad[: n[0], : n[0]] = X
          return X_pad
  
      A_n = A.shape[0]
      B_n = B.shape[0]
      n = np.sort([A_n, B_n])
      if method == "adopted":
          A = 2 * A - np.ones((A_n, A_n))
          B = 2 * B - np.ones((B_n, B_n))
  
      S_pad = np.zeros((n[1], n[1]))
      S_pad[:A_n, :B_n] = S
  
      if A.shape[0] == n[0]:
          A = pad(A, n)
      else:
          B = pad(B, n)
  
      return A, B, S_pad
  