content: import numbers
  
  import numpy as np
  
  from scipy import sparse
  from sklearn.cluster import KMeans
  
  from . import GraphCRF
  from ..inference import inference_dispatch
  
  
  def kmeans_init(X, Y, all_edges, n_labels, n_states_per_label,
                  symmetric=True):
      all_feats = []
      # iterate over samples
      for x, y, edges in zip(X, Y, all_edges):
          # first, get neighbor counts from nodes
          n_nodes = x.shape[0]
          labels_one_hot = np.zeros((n_nodes, n_labels), dtype=np.int)
          y = y.ravel()
          gx = np.ogrid[:n_nodes]
          labels_one_hot[gx, y] = 1
  
          size = np.prod(y.shape)
          graphs = [sparse.coo_matrix((np.ones(e.shape[0]), e.T), (size, size))
                    for e in edges]
          if symmetric:
              directions = [g + g.T for g in graphs]
          else:
              directions = [T for g in graphs for T in [g, g.T]]
          neighbors = [s * labels_one_hot.reshape(size, -1) for s in directions]
          neighbors = np.hstack(neighbors)
          # normalize (for borders)
          neighbors /= np.maximum(neighbors.sum(axis=1)[:, np.newaxis], 1)
  
          # add unaries
          features = np.hstack([x, neighbors])
          all_feats.append(features)
      all_feats_stacked = np.vstack(all_feats)
      Y_stacked = np.hstack(Y).ravel()
      # for each state, run k-means over whole dataset
      H = [np.zeros(y.shape, dtype=np.int) for y in Y]
      label_indices = np.hstack([0, np.cumsum(n_states_per_label)])
      for label in np.unique(Y_stacked):
          try:
              km = KMeans(n_clusters=n_states_per_label[label])
          except TypeError:
              # for old versions :-/
              km = KMeans(k=n_states_per_label[label])
          indicator = Y_stacked == label
          f = all_feats_stacked[indicator]
          km.fit(f)
          for feats_sample, y, h in zip(all_feats, Y, H):
              indicator_sample = y.ravel() == label
              if np.any(indicator_sample):
                  pred = km.predict(feats_sample[indicator_sample]).astype(np.int)
                  h.ravel()[indicator_sample] = pred + label_indices[label]
      return H
  
  
  class LatentGraphCRF(GraphCRF):
      """CRF with latent states for variables.
  
      This is also called "hidden dynamics CRF".
      For each output variable there is an additional variable which
      can encode additional states and interactions.
  
      Parameters
      ----------
      n_labels : int
          Number of states of output variables.
  
      n_featues : int or None (default=None).
          Number of input features per input variable.
          ``None`` means it is equal to ``n_labels``.
  
      n_states_per_label : int or list (default=2)
          Number of latent states associated with each observable state.
          Can be either an integer, which means the same number
          of hidden states will be used for all observable states, or a list
          of integers of length ``n_labels``.
  
      inference_method : string, default="ad3"
          Function to call to do inference and loss-augmented inference.
          Possible values are:
  
              - 'max-product' for max-product belief propagation.
                  Recommended for chains an trees. Loopy belief propagatin in case of a general graph.
              - 'lp' for Linear Programming relaxation using cvxopt.
              - 'ad3' for AD3 dual decomposition.
              - 'qpbo' for QPBO + alpha expansion.
              - 'ogm' for OpenGM inference algorithms.
      """
      def __init__(self, n_labels=None, n_features=None, n_states_per_label=2,
                   inference_method=None):
          self.n_labels = n_labels
          self.n_states_per_label = n_states_per_label
          GraphCRF.__init__(self, n_states=None, n_features=n_features,
                            inference_method=inference_method)
  
      def _set_size_joint_feature(self):
          if None in [self.n_features, self.n_labels]:
              return
  
          if isinstance(self.n_states_per_label, numbers.Integral):
              # same for all labels
              n_states_per_label = np.array([
                  self.n_states_per_label for i in range(self.n_labels)])
          else:
              n_states_per_label = np.array(self.n_states_per_label)
              if len(n_states_per_label) != self.n_labels:
                  raise ValueError("states_per_label must be integer"
                                   "or array-like of length n_labels. Got %s"
                                   % str(n_states_per_label))
          self.n_states_per_label = n_states_per_label
          self.n_states = np.sum(n_states_per_label)
  
          # compute mapping from latent states to labels
          ranges = np.cumsum(n_states_per_label)
          states_map = np.zeros(self.n_states, dtype=np.int)
          for l in range(1, self.n_labels):
              states_map[ranges[l - 1]: ranges[l]] = l
          self._states_map = states_map
          GraphCRF._set_size_joint_feature(self)
  
      def initialize(self, X, Y):
          n_features = X[0][0].shape[1]
          if self.n_features is None:
              self.n_features = n_features
          elif self.n_features != n_features:
              raise ValueError("Expected %d features, got %d"
                               % (self.n_features, n_features))
  
          n_labels = len(np.unique(np.hstack([y.ravel() for y in Y])))
          if self.n_labels is None:
              self.n_labels = n_labels
          elif self.n_labels != n_labels:
              raise ValueError("Expected %d states, got %d"
                               % (self.n_labels, n_labels))
          self._set_size_joint_feature()
          self._set_class_weight()
  
      def label_from_latent(self, h):
          return self._states_map[h]
  
      def init_latent(self, X, Y):
          # treat all edges the same
          edges = [[self._get_edges(x)] for x in X]
          features = np.array([self._get_features(x) for x in X])
          return kmeans_init(features, Y, edges, n_labels=self.n_labels,
                             n_states_per_label=self.n_states_per_label)
  
      def loss_augmented_inference(self, x, h, w, relaxed=False,
                                   return_energy=False):
          self.inference_calls += 1
          self._check_size_w(w)
          unary_potentials = self._get_unary_potentials(x, w)
          pairwise_potentials = self._get_pairwise_potentials(x, w)
          edges = self._get_edges(x)
          # do loss-augmentation
          for l in np.arange(self.n_states):
              # for each class, decrement features
              # for loss-agumention
              unary_potentials[self.label_from_latent(h)
                               != self.label_from_latent(l), l] += 1.
  
          return inference_dispatch(unary_potentials, pairwise_potentials, edges,
                                    self.inference_method, relaxed=relaxed,
                                    return_energy=return_energy)
  
      def latent(self, x, y, w):
          unary_potentials = self._get_unary_potentials(x, w)
          # forbid h that is incompoatible with y
          # by modifying unary params
          other_states = self._states_map != y[:, np.newaxis]
          max_entry = np.maximum(np.max(unary_potentials), 1)
          unary_potentials[other_states] = -1e2 * max_entry
          pairwise_potentials = self._get_pairwise_potentials(x, w)
          edges = self._get_edges(x)
          h = inference_dispatch(unary_potentials, pairwise_potentials, edges,
                                 self.inference_method, relaxed=False)
          if (self.label_from_latent(h) != y).any():
              print("inconsistent h and y")
              h = np.hstack([0, np.cumsum(self.n_states_per_label)])[y]
          return h
  
      def loss(self, h, h_hat):
          if isinstance(h_hat, tuple):
              return self.continuous_loss(h, h_hat[0])
          return GraphCRF.loss(self, self.label_from_latent(h),
                               self.label_from_latent(h_hat))
  
      def continuous_loss(self, y, y_hat):
          # continuous version of the loss
          # y_hat is the result of linear programming
          y_hat_org = np.zeros((y_hat.shape[0], self.n_labels))
          for s in range(self.n_states):
              y_hat_org[:, self._states_map[s]] += y_hat[:, s]
          y_org = self.label_from_latent(y)
          return GraphCRF.continuous_loss(self, y_org, y_hat_org)
  
      def base_loss(self, y, y_hat):
          if isinstance(y_hat, tuple):
              return GraphCRF.continuous_loss(self, y, y_hat)
          return GraphCRF.loss(self, y, y_hat)
  