// rnnlm/sampling-lm-test.cc
  
  // Copyright 2017  Ke Li
  //           2017  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.
  
  #include "rnnlm/sampling-lm.h"
  
  namespace kaldi {
  namespace rnnlm {
  
  class SamplingLmTest {
   public:
    typedef SamplingLm::HistType HistType;
    typedef SamplingLm::WeightedHistType WeightedHistType;
  
    explicit SamplingLmTest(SamplingLm *arpa) {
      arpa_ = arpa;
    }
    // This function reads in a list of histories and their weights from a file
    // only text form is supported
    void ReadHistories(std::istream &is, bool binary,
        WeightedHistType *histories);
  
    void TestUnigramDistribution();
  
    void TestGetDistribution(const WeightedHistType &histories);
   private:
    // This SamplingLm object is used to get accesses to private and protected
    // members in SamplingLm class
    SamplingLm *arpa_;
  };
  
  void SamplingLmTest::ReadHistories(std::istream &is, bool binary,
      WeightedHistType *histories) {
    if (binary) {
      KALDI_ERR << "binary-mode reading is not implemented for ArpaFileParser";
    }
    const fst::SymbolTable* sym = arpa_->Symbols();
    std::string line;
    KALDI_LOG << "Start reading histories from file...";
    // int32 ngram_order = arpa_->ngram_order_;
    int32 ngram_order = arpa_->Order();
    while (getline(is, line)) {
      std::istringstream is(line);
      std::istream_iterator<std::string> begin(is), end;
      std::vector<std::string> tokens(begin, end);
      HistType history;
      int32 word;
      BaseFloat hist_weight = 0;
      for (int32 i = 0; i < tokens.size() - 1; ++i) {
        word = sym->Find(tokens[i]);
        if (word == -1) { // fst::kNoSymbol
          KALDI_ERR << "Found history contains word that is not in Arpa LM";
        }
        history.push_back(word);
      }
      HistType h1;
      if (history.size() >= ngram_order) {
        HistType h(history.end() - ngram_order + 1, history.end());
        h1 = h;
      }
      if (!ConvertStringToReal(tokens.back(), &hist_weight)) {
        KALDI_ERR << arpa_->LineReference() << ": invalid history weight '"
          << tokens.back() << "'";
      }
      KALDI_ASSERT(hist_weight >= 0);
      std::pair<HistType, BaseFloat> hist_pair;
      if (history.size() >= ngram_order) {
        hist_pair = std::make_pair(h1, hist_weight);
      } else {
        hist_pair = std::make_pair(history, hist_weight);
      }
      (*histories).push_back(hist_pair);
    }
    KALDI_LOG << "Successfully reading histories from file.";
  }
  
  void SamplingLmTest::TestUnigramDistribution() {
    std::vector<BaseFloat> unigram_probs;
    unigram_probs = arpa_->GetUnigramDistribution();
    // Check 0 (epsilon) has probability 0.0
    KALDI_ASSERT(unigram_probs[0] == 0.0);
    // Assert the sum of unigram probs of all words is 1.0
    BaseFloat probsum = 0.0;
    for (int32 i = 0; i < unigram_probs.size(); ++i) {
      probsum += unigram_probs[i];
    }
    KALDI_ASSERT(ApproxEqual(probsum, 1.0));
  }
  
  void SamplingLmTest::TestGetDistribution(const WeightedHistType &histories) {
    // get total input weights of histories
    BaseFloat total_weights = 0.0;
    WeightedHistType::const_iterator it = histories.begin();
    for (; it != histories.end(); ++it) {
      total_weights += it->second;
    }
    BaseFloat unigram_weight = 0.0;
    BaseFloat non_unigram_probsum = 0.0;
    std::vector<std::pair<int32, BaseFloat> > pdf;
    unigram_weight = arpa_->GetDistribution(histories, &pdf);
    for (int32 i = 0; i < pdf.size(); ++i) {
      non_unigram_probsum += pdf[i].second;
    }
    // assert unigram weight plus total non_unigram probs equals
    // the total input histories' weights
    KALDI_ASSERT(ApproxEqual(unigram_weight + non_unigram_probsum, total_weights));
  }
  
  }  // namespace rnnlm
  }  // namespace kaldi
  
  int main(int argc, char **argv) {
    using namespace kaldi;
    using namespace kaldi::rnnlm;
  
    const char *usage = "";
    ParseOptions po(usage);
    po.Read(argc, argv);
    std::string arpa_file = "test/0.1k_3gram_unpruned.arpa",
      history_file = "test/hists";
  
    ArpaParseOptions options;
    fst::SymbolTable symbols;
  
    enum {
      kEps = 0,
      kBos, kEos, kUnk
    };
  
    // Use spaces on special symbols, so we rather fail than read them by mistake.
    symbols.AddSymbol(" <eps>", kEps);
    options.bos_symbol = symbols.AddSymbol("<s>", kBos);
    options.eos_symbol = symbols.AddSymbol("</s>", kEos);
    options.unk_symbol = symbols.AddSymbol("<unk>", kUnk);
    options.oov_handling = ArpaParseOptions::kAddToSymbols;
    SamplingLm arpa(options, &symbols);
  
    bool binary;
    Input k1(arpa_file);
    arpa.Read(k1.Stream());
  
    SamplingLmTest mdl(&arpa);
    mdl.TestUnigramDistribution();
  
    Input k2(history_file, &binary);
    SamplingLmTest::WeightedHistType histories;
    mdl.ReadHistories(k2.Stream(), binary, &histories);
    mdl.TestGetDistribution(histories);
    KALDI_LOG << "Tests for SamplingLm class succeed.";
    return 0;
  }