// lm/arpa-lm-compiler-test.cc
  
  // Copyright 2009-2011 Gilles Boulianne
  // Copyright 2016 Smart Action LLC (kkm)
  
  // 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.
  
  #include <iostream>
  #include <string>
  #include <sstream>
  
  #include "base/kaldi-error.h"
  #include "base/kaldi-math.h"
  #include "lm/arpa-lm-compiler.h"
  #include "util/kaldi-io.h"
  
  namespace kaldi {
  
  // Predefine some symbol values, because any integer is as good than any other.
  enum {
    kEps = 0,
    kDisambig,
    kBos, kEos,
  };
  
  // Number of random sentences for coverage test.
  static const int kRandomSentences = 50;
  
  // Creates an FST that generates any sequence of symbols taken from given
  // symbol table. The FST is then associated with the symbol table.
  static fst::StdVectorFst* CreateGenFst(bool seps, const fst::SymbolTable* pst) {
    fst::StdVectorFst* genFst = new fst::StdVectorFst;
    genFst->SetInputSymbols(pst);
    genFst->SetOutputSymbols(pst);
  
    fst::StdArc::StateId midId   = genFst->AddState();
    if (!seps) {
      fst::StdArc::StateId initId  = genFst->AddState();
      fst::StdArc::StateId finalId = genFst->AddState();
      genFst->SetStart(initId);
      genFst->SetFinal(finalId, fst::StdArc::Weight::One());
      genFst->AddArc(initId, fst::StdArc(kBos, kBos, 0, midId));
      genFst->AddArc(midId,  fst::StdArc(kEos, kEos, 0, finalId));
    } else {
      genFst->SetStart(midId);
      genFst->SetFinal(midId, fst::StdArc::Weight::One());
    }
  
    // Add a loop for each symbol in the table except the four special ones.
    fst::SymbolTableIterator si(*pst);
    for (si.Reset(); !si.Done(); si.Next()) {
      if (si.Value() == kBos || si.Value() == kEos ||
          si.Value() == kEps || si.Value() == kDisambig)
        continue;
      genFst->AddArc(midId, fst::StdArc(si.Value(), si.Value(),
                                        fst::StdArc::Weight::One(), midId));
    }
    return genFst;
  }
  
  // Compile given ARPA file.
  ArpaLmCompiler* Compile(bool seps, const string &infile) {
    ArpaParseOptions options;
    fst::SymbolTable symbols;
    // Use spaces on special symbols, so we rather fail than read them by mistake.
    symbols.AddSymbol(" <eps>", kEps);
    symbols.AddSymbol(" #0", kDisambig);
    options.bos_symbol = symbols.AddSymbol("<s>", kBos);
    options.eos_symbol = symbols.AddSymbol("</s>", kEos);
    options.oov_handling = ArpaParseOptions::kAddToSymbols;
  
    // Tests in this form cannot be run with epsilon substitution, unless every
    // random path is also fitted with a #0-transducing self-loop.
    ArpaLmCompiler* lm_compiler =
        new ArpaLmCompiler(options,
                           seps ? kDisambig : 0,
                           &symbols);
    {
      Input ki(infile);
      lm_compiler->Read(ki.Stream());
    }
    return lm_compiler;
  }
  
  // Add a state to an FSA after last_state, add a form last_state to the new
  // atate, and return the new state.
  fst::StdArc::StateId AddToChainFsa(fst::StdMutableFst* fst,
                                     fst::StdArc::StateId last_state,
                                     int64 symbol) {
    fst::StdArc::StateId next_state  = fst->AddState();
    fst->AddArc(last_state, fst::StdArc(symbol, symbol, 0, next_state));
    return next_state;
  }
  
  // Add a disambiguator-generating self loop to every state of an FST.
  void AddSelfLoops(fst::StdMutableFst* fst) {
    for (fst::StateIterator<fst::StdMutableFst> siter(*fst);
         !siter.Done(); siter.Next()) {
      fst->AddArc(siter.Value(),
                  fst::StdArc(kEps, kDisambig, 0, siter.Value()));
    }
  }
  
  // Compiles infile and then runs kRandomSentences random coverage tests on the
  // compiled FST.
  bool CoverageTest(bool seps, const string &infile) {
    // Compile ARPA model.
    ArpaLmCompiler* lm_compiler = Compile(seps, infile);
  
    // Create an FST that generates any sequence of symbols taken from the model
    // output.
    fst::StdVectorFst* genFst =
        CreateGenFst(seps, lm_compiler->Fst().OutputSymbols());
  
    int num_successes = 0;
    for (int32 i = 0; i < kRandomSentences; ++i) {
      // Generate a random sentence FST.
      fst::StdVectorFst sentence;
      RandGen(*genFst, &sentence);
      if (seps)
        AddSelfLoops(&sentence);
  
      fst::ArcSort(lm_compiler->MutableFst(), fst::StdOLabelCompare());
  
      // The past must successfullycompose with the LM FST.
      fst::StdVectorFst composition;
      Compose(sentence, lm_compiler->Fst(), &composition);
      if (composition.Start() != fst::kNoStateId)
        ++num_successes;
    }
  
    delete genFst;
    delete lm_compiler;
  
    bool ok = num_successes == kRandomSentences;
    if (!ok) {
      KALDI_WARN << "Coverage test failed on " << infile << ": composed "
                 << num_successes << "/" << kRandomSentences;
    }
    return ok;
  }
  
  bool ScoringTest(bool seps, const string &infile, const string& sentence,
                   float expected) {
    ArpaLmCompiler* lm_compiler = Compile(seps, infile);
    const fst::SymbolTable* symbols = lm_compiler->Fst().InputSymbols();
  
    // Create a sentence FST for scoring.
    fst::StdVectorFst sentFst;
    fst::StdArc::StateId state = sentFst.AddState();
    sentFst.SetStart(state);
    if (!seps) {
      state = AddToChainFsa(&sentFst, state, kBos);
    }
    std::stringstream ss(sentence);
    string word;
    while (ss >> word) {
      int64 word_sym = symbols->Find(word);
      KALDI_ASSERT(word_sym != -1);
      state = AddToChainFsa(&sentFst, state, word_sym);
    }
    if (!seps) {
      state = AddToChainFsa(&sentFst, state, kEos);
    }
    if (seps) {
      AddSelfLoops(&sentFst);
    }
    sentFst.SetFinal(state, 0);
    sentFst.SetOutputSymbols(symbols);
  
    // Do the composition and extract final weight.
    fst::StdVectorFst composed;
    fst::Compose(sentFst, lm_compiler->Fst(), &composed);
    delete lm_compiler;
  
    if (composed.Start() == fst::kNoStateId) {
      KALDI_WARN << "Test sentence " << sentence << " did not compose "
                 << "with the language model FST
  ";
      return false;
    }
  
    std::vector<fst::StdArc::Weight> shortest;
    fst::ShortestDistance(composed, &shortest, true);
    float actual = shortest[composed.Start()].Value();
  
    bool ok = ApproxEqual(expected, actual);
    if (!ok) {
      KALDI_WARN << "Scored " << sentence << " in " << infile
                 << ": Expected=" << expected << " actual=" << actual;
    }
    return ok;
  }
  
  bool ThrowsExceptionTest(bool seps, const string &infile) {
    try {
      // Make memory cleanup easy in both cases of try-catch block.
      std::unique_ptr<ArpaLmCompiler> compiler(Compile(seps, infile));
      return false;
    } catch (const KaldiFatalError&) {
      return true;
    }
  }
  
  }  // namespace kaldi
  
  bool RunAllTests(bool seps) {
    bool ok = true;
    ok &= kaldi::CoverageTest(seps, "test_data/missing_backoffs.arpa");
    ok &= kaldi::CoverageTest(seps, "test_data/unused_backoffs.arpa");
    ok &= kaldi::CoverageTest(seps, "test_data/input.arpa");
  
    ok &= kaldi::ScoringTest(seps, "test_data/input.arpa", "b b b a", 59.2649);
    ok &= kaldi::ScoringTest(seps, "test_data/input.arpa", "a b", 4.36082);
  
    ok &= kaldi::ThrowsExceptionTest(seps, "test_data/missing_bos.arpa");
  
    if (!ok) {
      KALDI_WARN << "Tests " << (seps ? "with" : "without")
                 << " epsilon substitution FAILED";
    }
    return ok;
  }
  
  int main(int argc, char *argv[]) {
    bool ok = true;
  
    ok &= RunAllTests(false);  // Without disambiguators (old behavior).
    ok &= RunAllTests(true);   // With epsilon substitution (new behavior).
  
    if (ok) {
      KALDI_LOG << "All tests passed";
      return 0;
    } else {
      KALDI_WARN << "Test FAILED";
      return 1;
    }
  }