// lat/word-align-lattice-lexicon.h

// Copyright 2013 Johns Hopkins University (Author: Daniel Povey)

// See ../../COPYING 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#ifndef KALDI_LAT_WORD_ALIGN_LATTICE_LEXICON_H_
#define KALDI_LAT_WORD_ALIGN_LATTICE_LEXICON_H_
#include <fst/fstlib.h>
#include <fst/fst-decl.h>

#include "base/kaldi-common.h"
#include "util/common-utils.h"
#include "fstext/fstext-lib.h"
#include "hmm/transition-model.h"
#include "lat/kaldi-lattice.h"

namespace kaldi {

/** Read the lexicon in the special format required for word alignment.  Each line has
   a series of integers on it (at least two on each line), representing:

   <old-word-id> <new-word-id> [<phone-id-1> [<phone-id-2> ... ] ]

   Here, <old-word-id> is the word-id that appears in the lattice before alignment, and
   <new-word-id> is the word-is that should appear in the lattice after alignment.  This
   is mainly useful when the lattice may have no symbol for the optional-silence arcs
   (so <old-word-id> would equal zero), but we want it to be output with a symbol on those
   arcs (so <new-word-id> would be nonzero).
   If the silence should not be added to the lattice, both <old-word-id> and <new-word-id>
   may be zero.

   This function is very simple: it just reads in a series of lines from a text file,
   each with at least two integers on them.
*/
bool ReadLexiconForWordAlign (std::istream &is,
                              std::vector<std::vector<int32> > *lexicon);



/// This class extracts some information from the lexicon and stores it
/// in a suitable form for the word-alignment code to use.
class WordAlignLatticeLexiconInfo {
 public:
  WordAlignLatticeLexiconInfo(const std::vector<std::vector<int32> > &lexicon);

  /// Returns true if this lexicon-entry can appear, intepreted as
  /// (output-word phone1 phone2 ...).  This is just used in testing code.
  bool IsValidEntry(const std::vector<int32> &entry) const;

  /// Purely for the testing code, we map words into equivalence classes derived
  /// from the mappings in the first two fields of each line in the lexicon.  This
  /// function maps from each word-id to the lowest member of its equivalence class.
  int32 EquivalenceClassOf(int32 word) const;
 protected:
  friend class LatticeLexiconWordAligner;

  void UpdateViabilityMap(const std::vector<int32> &lexicon_entry);
  void UpdateLexiconMap(const std::vector<int32> &lexicon_entry);
  void UpdateNumPhonesMap(const std::vector<int32> &lexicon_entry);
  void UpdateEquivalenceMap(const std::vector<std::vector<int32> > &lexicon);

  void FinalizeViabilityMap(); // sorts the vectors.

  /// The type ViabilityMap maps from sequences of phones (excluding the empty
  /// sequence), to the sets of all word-labels [on the input lattice] that
  /// could correspond to phone sequences that start with s [but are longer than
  /// s].  The sets of word-labels are represented as sorted vectors of int32
  /// Note: the zero word-label is included here.  This is used in a kind
  /// of co-accessibility test, to see whether it is worth extending this state
  /// by traversing arcs in the input lattice.
  typedef unordered_map<std::vector<int32>,
                        std::vector<int32>,
                        VectorHasher<int32> > ViabilityMap;

  /// This is a map from a vector (orig-word-symbol phone1 phone2 ... ) to
  /// the new word-symbol.  [todo: make sure the new word-symbol is always nonzero.]
  typedef unordered_map<std::vector<int32>, int32,
                        VectorHasher<int32> > LexiconMap;

  /// This is a map from the word-id (as present in the original lattice)
  /// to the minimum and maximum #phones of lexicon entries for that word.
  /// It helps improve efficiency.
  typedef unordered_map<int32, std::pair<int32, int32> > NumPhonesMap;

  /// This is used only in testing code; it defines a mapping from a word
  /// to the primary member of that word's equivalence-class.
  typedef unordered_map<int32, int32> EquivalenceMap;

  // The following three variables represent various types of information
  // gathered from the lexicon.
  LexiconMap lexicon_map_;
  NumPhonesMap num_phones_map_;
  ViabilityMap viability_map_;

  // As lexicon_map but in reverse sense w.r.t. words [we only
  // do this for asymmetric entries.]  Used only in testing code.
  LexiconMap reverse_lexicon_map_;

  // This is used only in testing code; it defines a mapping from a word
  // to the primary member of that word's equivalence-class.  If an index
  // is not present in the map, it's assumed to map to itself.
  EquivalenceMap equivalence_map_;
};


struct WordAlignLatticeLexiconOpts {
  int32 partial_word_label;
  bool reorder;
  bool test;
  bool allow_duplicate_paths;
  BaseFloat max_expand;

  WordAlignLatticeLexiconOpts(): partial_word_label(0), reorder(true),
                                 test(false), allow_duplicate_paths(false),
                                 max_expand(-1.0) { }

  void Register(OptionsItf *opts) {
    opts->Register("partial-word-label", &partial_word_label, "Numeric id of "
                   "word symbol that is to be used for arcs in the word-aligned "
                   "lattice corresponding to partial words at the end of "
                   "\"forced-out\" utterances (zero is OK)");
    opts->Register("reorder", &reorder, "True if the lattices were generated "
                   "from graphs that had the --reorder option true, relating to "
                   "reordering self-loops (typically true)");
    opts->Register("test", &test, "If true, testing code will be activated "
                   "(the purpose of this is to validate the algorithm).");
    opts->Register("allow-duplicate-paths", &allow_duplicate_paths, "Only "
                   "has an effect if --test=true.  If true, does not die "
                   "(only prints warnings) if duplicate paths are found. "
                   "This should only happen with very pathological lexicons, "
                   "e.g. as encountered in testing code.");
    opts->Register("max-expand", &max_expand, "If >0.0, the maximum ratio "
                   "by which we allow the lattice-alignment code to increase the #states"
                   "in a lattice (vs. the phone-aligned lattice) before we fail and "
                   "refuse to align the lattice.  This is helpful in order to "
                   "prevent 'pathological' lattices from causing the program to "
                   "exhaust memory.  Actual max-states is 1000 + max-expand * "
                   "orig-num-states.");
  }
};


/// Align lattice so that each arc has the transition-ids on it
/// that correspond to the word that is on that arc.  [May also have
/// epsilon arcs for optional silences.]
/// Returns true if everything was OK, false if there was any kind of
/// error including when the the lattice seems to have been "forced out"
/// (did not reach end state, resulting in partial word at end).
bool WordAlignLatticeLexicon(const CompactLattice &lat,
                             const TransitionModel &tmodel,
                             const WordAlignLatticeLexiconInfo &lexicon_info,
                             const WordAlignLatticeLexiconOpts &opts,
                             CompactLattice *lat_out);


/// This function is designed to crash if something went wrong with the
/// word-alignment of the lattice.  If was_ok==true (was_ok is the return status
/// of WordAlignLattice), it tests that, after removing any silence and
/// partial-word labels that may have been inserted by WordAlignLattice,
/// the word-aligned lattice is equivalent to the input.  It also verifies
/// that arcs are of 4 types:
///   properly-aligned word arcs, with a word label.
///   partial-word arcs, with the partial-word label.
///   silence arcs, with the silence label.
void TestWordAlignedLatticeLexicon(const CompactLattice &lat,
                                   const TransitionModel &tmodel,
                                   const std::vector<std::vector<int32> > &lexicon,
                                   const CompactLattice &aligned_lat,
                                   bool allow_duplicate_paths);

} // end namespace kaldi
#endif