text: # Copyright      2020  Mobvoi Inc.        (authors: Fangjun Kuang)
  #                      Xiaomi Corporation (authors: Haowen Qiu)
  #
  # See ../../../LICENSE for clarification regarding multiple authors
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from typing import Optional
  from typing import Tuple
  from typing import Union
  
  import torch
  
  from .fsa import Fsa
  from .ops import index
  from .ops import index_ragged
  from .ops import index_select
  from .symbol_table import SymbolTable
  import k2
  import k2.ragged
  import _k2
  
  
  def to_str(fsa: Fsa,
             openfst: bool = False) -> str:
      '''Convert an Fsa to a string.  This version prints out all integer
      labels and integer ragged labels on the same line as each arc, the
      same format accepted by Fsa.from_str().
  
      Note:
        The returned string can be used to construct an Fsa with Fsa.from_str(),
        but you would need to know the names of the auxiliary labels and ragged
        labels.
  
      Args:
        openfst:
          Optional. If true, we negate the scores during the conversion.
  
      Returns:
        A string representation of the Fsa.
      '''
      assert fsa.arcs.num_axes() == 2
      extra_labels = []
      ragged_labels = []
      for name, value in sorted(fsa.named_tensor_attr(include_scores=False)):
          if isinstance(value, torch.Tensor) and value.dtype == torch.int32:
              extra_labels.append(value)
          elif isinstance(value, _k2.RaggedInt):
              ragged_labels.append(value)
  
      return _k2.fsa_to_str(fsa.arcs, openfst=openfst,
                            extra_labels=extra_labels,
                            ragged_labels=ragged_labels)
  
  
  def to_str_simple(fsa: Fsa,
                    openfst: bool = False) -> str:
      '''Convert an Fsa to a string.  This is less complete than Fsa.to_str(),
      fsa.__str__(), or to_str_full(), meaning it prints only fsa.aux_labels and
      no ragged labels, not printing any other attributes.  This is used in
      testing.
  
      Note:
        The returned string can be used to construct an Fsa.  See also to_str().
  
      Args:
        openfst:
          Optional. If true, we negate the scores during the conversion.
  
      Returns:
        A string representation of the Fsa.
      '''
      assert fsa.arcs.num_axes() == 2
      if hasattr(fsa, 'aux_labels') and isinstance(fsa.aux_labels, torch.Tensor):
          aux_labels = [fsa.aux_labels.to(torch.int32)]
      else:
          aux_labels = []
      return _k2.fsa_to_str(fsa.arcs, openfst, aux_labels, [])
  
  
  def to_tensor(fsa: Fsa) -> torch.Tensor:
      '''Convert an Fsa to a Tensor.
  
      You can save the tensor to disk and read it later
      to construct an Fsa.
  
      Note:
        The returned Tensor contains only the transition rules, e.g.,
        arcs. You may want to save its aux_labels separately if any.
  
      Args:
        fsa:
          The input Fsa.
      Returns:
        A `torch.Tensor` of dtype `torch.int32`. It is a 2-D tensor
        if the input is a single FSA. It is a 1-D tensor if the input
        is a vector of FSAs.
      '''
      return _k2.fsa_to_tensor(fsa.arcs)
  
  
  def to_dot(fsa: Fsa, title: Optional[str] = None) -> 'Digraph':  # noqa
      '''Visualize an Fsa via graphviz.
  
      Note:
        Graphviz is needed only when this function is called.
  
      Args:
        fsa:
          The input FSA to be visualized.
  
        title:
          Optional. The title of the resulting visualization.
      Returns:
        a Diagraph from grahpviz.
      '''
  
      try:
          import graphviz
      except Exception:
          print(
              'You cannot use `to_dot` unless the graphviz package is installed.'
          )
          raise
  
      assert len(fsa.shape) == 2, 'FsaVec is not supported'
      if hasattr(fsa, 'aux_labels'):
          aux_labels = fsa.aux_labels
          name = 'WFST'
      else:
          aux_labels = None
          name = 'WFSA'
  
      def convert_aux_label_to_symbol(
              aux_labels: Union[torch.Tensor, _k2.RaggedInt],
              arc_index: int,
              symbols: Optional[SymbolTable] = None) -> str:
          '''Convert aux_label(s) to symbol(s).
  
          Args:
            aux_labels:
              The aux_labels of an FSA.
            arc_index:
              The index of the arc.
            symbols:
              Symbol table of the FSA associated with the `aux_labels`.
          Returns:
            If `aux_labels` is a torch.Tensor, it returns a single string.
            If `aux_labels` is a ragged tensor, it returns a string with symbols
            separated by a space.
          '''
          if isinstance(aux_labels, torch.Tensor):
              ans = int(aux_labels[arc_index])
              if ans != -1 and symbols is not None:
                  ans = symbols[ans]
              return f':{ans}'
          assert isinstance(aux_labels, _k2.RaggedInt)
          assert aux_labels.num_axes() == 2
          row_splits = aux_labels.row_splits(1).cpu()
          begin = row_splits[arc_index]
          end = row_splits[arc_index + 1]
          if end == begin:
              return ':<eps>'
  
          labels = aux_labels.values()[begin:end]
          ans = []
          for label in labels.tolist():
              if label == -1:
                  ans.append('-1')
              elif symbols is not None:
                  ans.append(symbols[label])
              else:
                  ans.append(f'{label}')
          return f':{" ".join(ans)}'
  
      graph_attr = {
          'rankdir': 'LR',
          'size': '8.5,11',
          'center': '1',
          'orientation': 'Portrait',
          'ranksep': '0.4',
          'nodesep': '0.25',
      }
      if title is not None:
          graph_attr['label'] = title
  
      default_node_attr = {
          'shape': 'circle',
          'style': 'bold',
          'fontsize': '14',
      }
  
      final_state_attr = {
          'shape': 'doublecircle',
          'style': 'bold',
          'fontsize': '14',
      }
  
      final_state = -1
      dot = graphviz.Digraph(name=name, graph_attr=graph_attr)
  
      seen = set()
      i = -1
      for arc, weight in zip(fsa.arcs.values()[:, :-1], fsa.scores.tolist()):
          i += 1
          src_state, dst_state, label = arc.tolist()
          src_state = str(src_state)
          dst_state = str(dst_state)
          label = int(label)
          if label == -1:
              final_state = dst_state
          if src_state not in seen:
              dot.node(src_state, label=src_state, **default_node_attr)
              seen.add(src_state)
  
          if dst_state not in seen:
              if dst_state == final_state:
                  dot.node(dst_state, label=dst_state, **final_state_attr)
  
              else:
                  dot.node(dst_state, label=dst_state, **default_node_attr)
              seen.add(dst_state)
          if aux_labels is not None:
              if hasattr(fsa, 'aux_labels_sym'):
                  aux_label = convert_aux_label_to_symbol(
                      aux_labels, i, fsa.aux_labels_sym)
              else:
                  aux_label = convert_aux_label_to_symbol(aux_labels, i, None)
              aux_label = aux_label.replace('<eps>', 'ε')
          else:
              aux_label = ''
  
          if hasattr(fsa, 'labels_sym') and label != -1:
              label = fsa.labels_sym.get(label)
              if label == '<eps>':
                  label = 'ε'
  
          weight = f'{weight:.2f}'.rstrip('0').rstrip('.')
          dot.edge(src_state, dst_state, label=f'{label}{aux_label}/{weight}')
      return dot
  
  
  def create_fsa_vec(fsas):
      '''Create an FsaVec from a list of FSAs
  
      We use the following rules to set the attributes of the output FsaVec:
  
      - For tensor attributes, we assume that all input FSAs have the same
        attribute name and the values are concatenated.
  
      - For non-tensor attributes, if any two of the input FSAs have the same
        attribute name, then we assume that their attribute values are equal and
        the output FSA will inherit the attribute.
  
      Args:
        fsas:
          A list of `Fsa`. Each element must be a single FSA.
  
      Returns:
        An instance of :class:`Fsa` that represents a FsaVec.
      '''
      ragged_arc_list = list()
      for fsa in fsas:
          assert len(fsa.shape) == 2
          ragged_arc_list.append(fsa.arcs)
  
      ragged_arcs = _k2.create_fsa_vec(ragged_arc_list)
      fsa_vec = Fsa(ragged_arcs)
  
      tensor_attr_names = set(
          name for name, _ in fsa.named_tensor_attr() for fsa in fsas)
      for name in tensor_attr_names:
          values = []
          for fsa in fsas:
              values.append(getattr(fsa, name))
          if isinstance(values[0], torch.Tensor):
              value = torch.cat(values)
          else:
              assert isinstance(values[0], _k2.RaggedInt)
              value = k2.ragged.cat(values, axis=0)
          setattr(fsa_vec, name, value)
  
      non_tensor_attr_names = set(
          name for name, _ in fsa.named_non_tensor_attr() for fsa in fsas)
  
      for name in non_tensor_attr_names:
          if name == 'properties':
              continue
  
          for fsa in fsas:
              value = getattr(fsa, name, None)
              if value is not None:
                  if hasattr(fsa_vec, name):
                      assert getattr(fsa_vec, name) == value
                  else:
                      setattr(fsa_vec, name, value)
      return fsa_vec
  
  
  def is_rand_equivalent(a: Fsa,
                         b: Fsa,
                         log_semiring: bool,
                         beam: float = float('inf'),
                         treat_epsilons_specially: bool = True,
                         delta: float = 1e-6,
                         npath: int = 100) -> bool:
      '''Check if the Fsa `a` appears to be equivalent to `b` by
      randomly checking some symbol sequences in them.
  
      Caution:
        It works only on CPU.
  
      Args:
        a:
          One of the input FSA. It can be either a single FSA or an FsaVec.
          Must be top-sorted and on CPU.
        b:
          The other input FSA. It must have the same NumAxes() as a.
          Must be top-sorted and on CPU.
        log_semiring:
          The semiring to be used for all weight measurements;
          if false then we use 'max' on alternative paths; if
          true we use 'log-add'.
        beam:
           beam > 0 that affects pruning; the algorithm will only check
           paths within `beam` of the total score of the lattice (for
           tropical semiring, it's max weight over all paths from start
           state to final state; for log semiring, it's log-sum probs over
           all paths) in `a` or `b`.
        treat_epsilons_specially:
           We'll do `intersection` between generated path and a or b when
           check equivalence. Generally, if it's true, we will treat
           epsilons as epsilon when doing intersection; Otherwise, epsilons
           will just be treated as any other symbol.
        delta:
           Tolerance for path weights to check the equivalence.
           If abs(weights_a, weights_b) <= delta, we say the two
           paths are equivalent.
        npath:
           The number of paths will be generated to check the
           equivalence of `a` and `b`
      Returns:
         True if the Fsa `a` appears to be equivalent to `b` by randomly
         generating `npath` paths from one of them and then checking if the symbol
         sequence exists in the other one and if the total weight for that symbol
         sequence is the same in both FSAs.
      '''
      return _k2.is_rand_equivalent(a.arcs, b.arcs, log_semiring, beam,
                                    treat_epsilons_specially, delta, npath)
  
  
  def create_sparse(rows: torch.Tensor,
                    cols: torch.Tensor,
                    values: torch.Tensor,
                    size: Optional[Tuple[int, int]] = None,
                    min_col_index: Optional[int] = None):
      '''This is a utility function that creates a (torch) sparse matrix likely
      intended to represent posteriors.  The likely usage is something like
      (for example)::
  
          post = k2.create_sparse(fsa.seqframe, fsa.phones,
                                  fsa.get_arc_post(True,True).exp(),
                                  min_col_index=1)
  
      (assuming `seqframe` and `phones` were integer-valued attributes of `fsa`).
  
      Args:
        rows:
          Row indexes of the sparse matrix (a torch.Tensor), which must have
          values >= 0; likely `fsa.seqframe`.   Must have row_indexes.dim == 1.
          Will be converted to `dtype=torch.long`
        cols:
          Column indexes of the sparse matrix, with the same shape as `rows`.
          Will be converted to `dtype=torch.long`
        values:
          Values of the sparse matrix, likely of dtype float or double, with
          the same shape as `rows` and `cols`.
        size:
          Optional. If not None, it is assumed to be a tuple containing
          `(num_frames, highest_phone_plus_one)`
        min_col_index:
          If provided, before the sparse tensor is constructed we will filter out
          elements with `cols[i] < min_col_index`.  Will likely be 0 or 1, if
          set.  This is necessary if `col_indexes` may have values less than 0,
          or if you want to filter out 0 values (e.g. as representing blanks).
  
      Returns:
        Returns a torch.Tensor that is sparse with coo (coordinate) format,
        i.e. `layout=torch.sparse_coo` (which is actually the only sparse format
        that torch currently supports).
      '''
      assert rows.ndim == cols.ndim == 1
      assert rows.numel() == cols.numel() == values.numel()
  
      if min_col_index is not None:
          assert isinstance(min_col_index, int)
          kept_indexes = cols >= min_col_index
          rows = rows[kept_indexes]
          cols = cols[kept_indexes]
          values = values[kept_indexes]
      if size is not None:
          return torch.sparse_coo_tensor(torch.stack([rows, cols]),
                                         values,
                                         size=size,
                                         device=values.device,
                                         requires_grad=values.requires_grad)
      else:
          return torch.sparse_coo_tensor(torch.stack([rows, cols]),
                                         values,
                                         device=values.device,
                                         requires_grad=values.requires_grad)
  
  
  def fsa_from_unary_function_tensor(src: Fsa, dest_arcs: _k2.RaggedArc,
                                     arc_map: torch.Tensor) -> Fsa:
      '''Create an Fsa object, including autograd logic and propagating
      properties from the source FSA.
  
      This is intended to be called from unary functions on FSAs where the arc_map
      is a Tensor of int32 (i.e. not ragged).
  
      Args:
        src:
          The source Fsa, i.e. the arg to the unary function.
        dest_arcs:
          The raw output of the unary function, as output by whatever C++
          algorithm we used.
        arc_map:
          A map from arcs in `dest_arcs` to the corresponding arc-index in `src`,
          or -1 if the arc had no source arc (e.g. added epsilon self-loops).
      Returns:
        Returns the resulting Fsa, with properties propagated appropriately, and
        autograd handled.
      '''
      dest = Fsa(dest_arcs)
  
      for name, value in src.named_tensor_attr(include_scores=False):
          if isinstance(value, torch.Tensor):
              filler = float(src.get_filler(name))
              setattr(dest, name, index_select(value, arc_map,
                                               default_value=filler))
          else:
              setattr(dest, name, index(value, arc_map))
  
      for name, value in src.named_non_tensor_attr():
          setattr(dest, name, value)
  
      k2.autograd_utils.phantom_index_select_scores(dest, src.scores, arc_map)
      return dest
  
  
  def fsa_from_unary_function_ragged(src: Fsa, dest_arcs: _k2.RaggedArc,
                                     arc_map: _k2.RaggedInt,
                                     remove_filler: bool = True) -> Fsa:
      '''Create an Fsa object, including autograd logic and propagating
      properties from the source FSA.
  
      This is intended to be called from unary functions on FSAs where the arc_map
      is an instance of _k2.RaggedInt.
  
      Args:
        src:
          The source Fsa, i.e. the arg to the unary function.
        dest_arcs:
          The raw output of the unary function, as output by whatever C++
          algorithm we used.
        arc_map:
          A map from arcs in `dest_arcs` to the corresponding arc-index in `src`,
          or -1 if the arc had no source arc (e.g. :func:`remove_epsilon`).
        remove_filler:
          If true, for each attribute that is linear in `src` and ragged
          in the result, after turning it into a ragged tensor we will
          remove all items that are equal to the filler for that attribute
          (0 by default; see Fsa.get_filler()).  Attribute values on final-arcs
          that are equal to -1 will also be treated as fillers and removed,
          if remove_filler==True.
      Returns:
          Returns the resulting Fsa, with properties propagated appropriately, and
          autograd handled.
      '''
      dest = Fsa(dest_arcs)
  
      for name, value in src.named_tensor_attr(include_scores=False):
          if remove_filler and isinstance(value, torch.Tensor) and \
             value.dtype == torch.int32:
              filler = src.get_filler(name)
              # when removing fillers for `aux_labels`, we need to treat -1 as a
              # filler where it is on a final-arc (i.e. turn it into the actual
              # filler, so it will be later removed by remove_values_eq).  We
              # assume that `dest` has been checked for validity, so the presence
              # of -1 as the label precisely indicates final-arcs.
              if filler != -1:
                  value = value.clone()
                  value[torch.where(torch.logical_and(
                      src.labels == -1, value == -1))] = filler
              new_value = index(value, arc_map)
              setattr(dest, name, k2.ragged.remove_values_eq(new_value, filler))
          else:
              new_value = index(value, arc_map)
              setattr(dest, name, new_value)
  
      for name, value in src.named_non_tensor_attr():
          setattr(dest, name, value)
  
      k2.autograd_utils.phantom_index_and_sum_scores(dest, src.scores, arc_map)
  
      return dest
  
  
  def fsa_from_binary_function_tensor(
          a_fsa: Fsa,
          b_fsa: Fsa,
          dest_arcs: _k2.RaggedArc,
          a_arc_map: torch.Tensor,
          b_arc_map: torch.Tensor) -> Fsa:
  
      '''Create an Fsa object, including autograd logic and propagating
      properties from the source FSAs.
  
      This is intended to be called from binary functions on FSAs where the
      arc_map is a Tensor of int32 (i.e. not ragged).
  
      Caution: Only the attributes with dtype `torch.float32` will be merged,
               other kinds of attributes with the same name are discarded.
  
      Args:
        a_fsa:
          The source Fsa, i.e. the arg to the binary function.
        b_fsa:
          The other source Fsa.
        dest_arcs:
          The raw output of the binary function, as output by whatever C++
          algorithm we used.
        a_arc_map:
          A map from arcs in `dest_arcs` to the corresponding arc-index in `a_fsa`
          or -1 if the arc had no source arc (e.g. added epsilon self-loops).
        a_arc_map:
          A map from arcs in `dest_arcs` to the corresponding arc-index in `b_fsa`
          or -1 if the arc had no source arc (e.g. added epsilon self-loops).
      Returns:
        Returns the resulting Fsa, with properties propagated appropriately, and
        autograd handled.
      '''
  
      out_fsa = Fsa(dest_arcs)
  
      for name, a_value in a_fsa.named_tensor_attr():
          # we include 'scores' in the attributes; this enables the
          # autograd to work.
          filler = float(a_fsa.get_filler(name))
          if hasattr(b_fsa, name):
              # Both a_fsa and b_fsa have this attribute.
              # We only support attributes with dtype `torch.float32`.
              # Other kinds of attributes are discarded.
              if a_value.dtype != torch.float32:
                  raise AttributeError("We don't support propagating two "
                                       "attributes with the same name that are "
                                       "not real-valued, in intersection: " +
                                       name)
              b_value = getattr(b_fsa, name)
              assert b_value.dtype == torch.float32
              # The following will actually overwrite `scores` with the same
              # value it had before; but this enables the autograd to work since
              # we do it using torch mechanisms.
              value = index_select(a_value, a_arc_map, default_value=filler) \
                      + index_select(b_value, b_arc_map, default_value=filler)
              setattr(out_fsa, name, value)
          else:
              # only a_fsa has this attribute, copy it via arc_map
              if isinstance(a_value, torch.Tensor):
                  value = index_select(a_value, a_arc_map, default_value=filler)
              else:
                  assert isinstance(a_value, _k2.RaggedInt)
                  value = index_ragged(a_value, a_arc_map)
              setattr(out_fsa, name, value)
  
      for name, b_value in b_fsa.named_tensor_attr():
          if not hasattr(out_fsa, name):
              if isinstance(b_value, torch.Tensor):
                  filler = float(b_fsa.get_filler(name))
                  value = index_select(b_value, b_arc_map, default_value=filler)
              else:
                  assert isinstance(b_value, _k2.RaggedInt)
                  value = index_ragged(b_value, b_arc_map)
              setattr(out_fsa, name, value)
  
      for name, a_value in a_fsa.named_non_tensor_attr():
          setattr(out_fsa, name, a_value)
  
      for name, b_value in b_fsa.named_non_tensor_attr():
          if not hasattr(out_fsa, name):
              setattr(out_fsa, name, b_value)
  
      return out_fsa
  
  
  def random_fsa(acyclic: bool = True,
                 max_symbol: int = 50,
                 min_num_arcs: int = 0,
                 max_num_arcs: int = 1000) -> Fsa:
  
      '''Generate a random Fsa.
      Args:
        acyclic:
          If true, generated Fsa will be acyclic.
        max_symbol:
          Maximum symbol on arcs. Generated arcs' symbols will be in range
          [-1,max_symbol], note -1 is kFinalSymbol; must be at least 0;
       min_num_arcs:
         Minimum number of arcs; must be at least 0.
       max_num_arcs:
         Maximum number of arcs; must be >= min_num_arcs.
      '''
  
      random_arcs = _k2.random_fsa(acyclic, max_symbol, min_num_arcs,
                                   max_num_arcs)
      return Fsa(random_arcs)
  
  
  def random_fsa_vec(min_num_fsas: int = 1,
                     max_num_fsas: int = 1000,
                     acyclic: bool = True,
                     max_symbol: int = 50,
                     min_num_arcs: int = 0,
                     max_num_arcs: int = 1000) -> Fsa:
  
      '''Generate a random FsaVec.
      Args:
        min_num_fsas:
          Minimum number of fsas we'll generated in the returned FsaVec;
          must be at least 1.
        max_num_fsas:
          Maximum number of fsas we'll generated in the returned FsaVec;
          must be >= min_num_fsas.
        acyclic:
          If true, generated Fsas will be acyclic.
        max_symbol:
          Maximum symbol on arcs. Generated arcs' symbols will be in range
          [-1,max_symbol], note -1 is kFinalSymbol; must be at least 0;
        min_num_arcs:
          Minimum number of arcs in each Fsa; must be at least 0.
        max_num_arcs:
          Maximum number of arcs in each Fsa; must be >= min_num_arcs.
      '''
  
      random_arcs = _k2.random_fsa_vec(min_num_fsas, max_num_fsas, acyclic,
                                       max_symbol, min_num_arcs, max_num_arcs)
      return Fsa(random_arcs)
  