// Copyright 2016 Google Inc.
//
// 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.!

#include "unigram_model.h"

#include <cmath>
#include <map>
#include <string>
#include <vector>

#include "sentencepiece_model.pb.h"
#include "sentencepiece_processor.h"
#include "testharness.h"
#include "third_party/absl/strings/str_cat.h"
#include "third_party/absl/strings/str_join.h"
#include "util.h"

namespace sentencepiece {
namespace unigram {

TEST(LatticeTest, SetSentenceTest) {
  Lattice lattice;

  EXPECT_EQ(0, lattice.size());
  EXPECT_EQ(0, lattice.utf8_size());

  lattice.SetSentence("");
  EXPECT_EQ(0, lattice.size());
  EXPECT_EQ(0, lattice.utf8_size());
  EXPECT_STREQ("", lattice.sentence());
  EXPECT_STREQ("", lattice.surface(0));

  lattice.SetSentence("test");
  EXPECT_EQ(4, lattice.size());
  EXPECT_EQ(4, lattice.utf8_size());
  EXPECT_STREQ("test", lattice.sentence());
  EXPECT_STREQ("test", lattice.surface(0));
  EXPECT_STREQ("est", lattice.surface(1));
  EXPECT_STREQ("st", lattice.surface(2));
  EXPECT_STREQ("t", lattice.surface(3));

  Lattice::Node *bos = lattice.bos_node();
  Lattice::Node *eos = lattice.eos_node();

  EXPECT_EQ(-1, bos->id);
  EXPECT_EQ(-1, eos->id);
  EXPECT_EQ(bos, lattice.end_nodes(0).front());
  EXPECT_EQ(eos, lattice.begin_nodes(4).front());

  lattice.SetSentence("テストab");
  EXPECT_EQ(5, lattice.size());
  EXPECT_EQ(11, lattice.utf8_size());
  EXPECT_STREQ("テストab", lattice.sentence());
  EXPECT_STREQ("テストab", lattice.surface(0));
  EXPECT_STREQ("ストab", lattice.surface(1));
  EXPECT_STREQ("トab", lattice.surface(2));
  EXPECT_STREQ("ab", lattice.surface(3));
  EXPECT_STREQ("b", lattice.surface(4));

  lattice.Clear();
  EXPECT_EQ(0, lattice.size());
  EXPECT_EQ(0, lattice.utf8_size());
}

TEST(LatticeTest, InsertTest) {
  Lattice lattice;
  lattice.SetSentence("ABあい");

  Lattice::Node *node[7];
  node[0] = lattice.Insert(0, 1);
  node[1] = lattice.Insert(1, 1);
  node[2] = lattice.Insert(2, 1);
  node[3] = lattice.Insert(3, 1);
  node[4] = lattice.Insert(0, 2);
  node[5] = lattice.Insert(1, 2);
  node[6] = lattice.Insert(2, 2);

  EXPECT_EQ("A", node[0]->piece);
  EXPECT_EQ("B", node[1]->piece);
  EXPECT_EQ("あ", node[2]->piece);
  EXPECT_EQ("い", node[3]->piece);
  EXPECT_EQ("AB", node[4]->piece);
  EXPECT_EQ("Bあ", node[5]->piece);
  EXPECT_EQ("あい", node[6]->piece);

  EXPECT_EQ("A", node[0]->piece);
  EXPECT_EQ("B", node[1]->piece);
  EXPECT_EQ("あ", node[2]->piece);
  EXPECT_EQ("い", node[3]->piece);
  EXPECT_EQ("AB", node[4]->piece);
  EXPECT_EQ("Bあ", node[5]->piece);
  EXPECT_EQ("あい", node[6]->piece);

  EXPECT_EQ(0, node[0]->pos);
  EXPECT_EQ(1, node[1]->pos);
  EXPECT_EQ(2, node[2]->pos);
  EXPECT_EQ(3, node[3]->pos);
  EXPECT_EQ(0, node[4]->pos);
  EXPECT_EQ(1, node[5]->pos);
  EXPECT_EQ(2, node[6]->pos);

  EXPECT_EQ(1, node[0]->length);
  EXPECT_EQ(1, node[1]->length);
  EXPECT_EQ(1, node[2]->length);
  EXPECT_EQ(1, node[3]->length);
  EXPECT_EQ(2, node[4]->length);
  EXPECT_EQ(2, node[5]->length);
  EXPECT_EQ(2, node[6]->length);

  EXPECT_EQ(0, lattice.bos_node()->node_id);
  EXPECT_EQ(1, lattice.eos_node()->node_id);
  EXPECT_EQ(2, node[0]->node_id);
  EXPECT_EQ(3, node[1]->node_id);
  EXPECT_EQ(4, node[2]->node_id);
  EXPECT_EQ(5, node[3]->node_id);
  EXPECT_EQ(6, node[4]->node_id);
  EXPECT_EQ(7, node[5]->node_id);
  EXPECT_EQ(8, node[6]->node_id);

  EXPECT_EQ(2, lattice.begin_nodes(0).size());
  EXPECT_EQ(2, lattice.begin_nodes(1).size());
  EXPECT_EQ(2, lattice.begin_nodes(2).size());
  EXPECT_EQ(1, lattice.begin_nodes(3).size());
  EXPECT_EQ(1, lattice.begin_nodes(4).size());  // EOS

  EXPECT_EQ(1, lattice.end_nodes(0).size());  // BOS
  EXPECT_EQ(1, lattice.end_nodes(1).size());
  EXPECT_EQ(2, lattice.end_nodes(2).size());
  EXPECT_EQ(2, lattice.end_nodes(3).size());
  EXPECT_EQ(2, lattice.end_nodes(4).size());

  EXPECT_EQ(node[0], lattice.begin_nodes(0)[0]);
  EXPECT_EQ(node[4], lattice.begin_nodes(0)[1]);
  EXPECT_EQ(node[1], lattice.begin_nodes(1)[0]);
  EXPECT_EQ(node[5], lattice.begin_nodes(1)[1]);
  EXPECT_EQ(node[2], lattice.begin_nodes(2)[0]);
  EXPECT_EQ(node[6], lattice.begin_nodes(2)[1]);
  EXPECT_EQ(node[3], lattice.begin_nodes(3)[0]);
  EXPECT_EQ(lattice.eos_node(), lattice.begin_nodes(4)[0]);

  EXPECT_EQ(lattice.bos_node(), lattice.end_nodes(0)[0]);
  EXPECT_EQ(node[0], lattice.end_nodes(1)[0]);
  EXPECT_EQ(node[1], lattice.end_nodes(2)[0]);
  EXPECT_EQ(node[4], lattice.end_nodes(2)[1]);
  EXPECT_EQ(node[2], lattice.end_nodes(3)[0]);
  EXPECT_EQ(node[5], lattice.end_nodes(3)[1]);
  EXPECT_EQ(node[3], lattice.end_nodes(4)[0]);
  EXPECT_EQ(node[6], lattice.end_nodes(4)[1]);
}

TEST(LatticeTest, ViterbiFromIncompleteLatticeTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");
  EXPECT_TRUE(lattice.Viterbi().first.empty());

  // Still incomplete
  lattice.Insert(0, 1);
  EXPECT_TRUE(lattice.Viterbi().first.empty());

  lattice.Insert(1, 1);
  lattice.Insert(2, 1);
  lattice.Viterbi();
}

std::string GetTokenized(const std::vector<Lattice::Node *> &nodes) {
  std::vector<std::string> tokens;
  for (auto *node : nodes) {
    tokens.push_back(std::string(node->piece));
  }
  return absl::StrJoin(tokens, " ");
}

void InsertWithScore(Lattice *lattice, int pos, int length, float score) {
  lattice->Insert(pos, length)->score = score;
}

void InsertWithScoreAndId(Lattice *lattice, int pos, int length, float score,
                          int id) {
  auto *node = lattice->Insert(pos, length);
  node->score = score;
  node->id = id;
}

TEST(LatticeTest, ViterbiTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScore(&lattice, 0, 1, 0.0);  // A
  InsertWithScore(&lattice, 1, 1, 0.0);  // B
  InsertWithScore(&lattice, 2, 1, 0.0);  // C
  EXPECT_EQ("A B C", GetTokenized(lattice.Viterbi().first));

  InsertWithScore(&lattice, 0, 2, 2.0);  // AB
  EXPECT_EQ("AB C", GetTokenized(lattice.Viterbi().first));

  InsertWithScore(&lattice, 1, 2, 5.0);  // BC
  EXPECT_EQ("A BC", GetTokenized(lattice.Viterbi().first));

  InsertWithScore(&lattice, 0, 3, 10.0);  // ABC
  EXPECT_EQ("ABC", GetTokenized(lattice.Viterbi().first));
}

TEST(LatticeTest, NBestTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScore(&lattice, 0, 1, 0.0);   // A
  InsertWithScore(&lattice, 1, 1, 0.0);   // B
  InsertWithScore(&lattice, 2, 1, 0.0);   // C
  InsertWithScore(&lattice, 0, 2, 2.0);   // AB
  InsertWithScore(&lattice, 1, 2, 5.0);   // BC
  InsertWithScore(&lattice, 0, 3, 10.0);  // ABC

  auto nbests = lattice.NBest(10, false, 0.0);
  EXPECT_EQ(4, nbests.size());

  EXPECT_EQ("ABC", GetTokenized(nbests[0].first));
  EXPECT_EQ("A BC", GetTokenized(nbests[1].first));
  EXPECT_EQ("AB C", GetTokenized(nbests[2].first));
  EXPECT_EQ("A B C", GetTokenized(nbests[3].first));

  auto nbests0 = lattice.NBest(0, false, 0.0);
  EXPECT_TRUE(nbests0.empty());

  auto nbests1 = lattice.NBest(1, false, 0.0);
  EXPECT_EQ(nbests1.size(), 1);
}

TEST(LatticeTest, NBestSampleTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScore(&lattice, 0, 1, 0.0);  // A
  InsertWithScore(&lattice, 1, 1, 0.0);  // B
  InsertWithScore(&lattice, 2, 1, 0.1);  // C
  InsertWithScore(&lattice, 0, 2, 0.2);  // AB
  InsertWithScore(&lattice, 1, 2, 0.5);  // BC
  InsertWithScore(&lattice, 0, 3, 1.0);  // ABC

  // Calculate expected probabilities of each path
  // Note that sampling without replacement affects the expected frequencies!
  const std::vector<double> kTheta = {0.0, 0.01, 0.5, 0.7, 1.0};
  for (const auto theta : kTheta) {
    std::vector<std::string> strings = {"ABC", "AB C", "A BC", "A B C"};
    std::map<std::string, float> probs;
    probs["ABC"] = std::exp(theta * 1.0);
    probs["AB C"] = std::exp(theta * (0.2 + 0.1));
    probs["A BC"] = std::exp(theta * (0.0 + 0.5));
    probs["A B C"] = std::exp(theta * (0.0 + 0.0 + 0.1));

    for (const auto &it : strings) {
      EXPECT_EQ(1, probs.count(it));
    }

    double Z = 0.0;
    for (const auto &it : probs) Z += it.second;
    for (auto &it : probs) it.second /= Z;

    std::map<std::pair<std::string, std::string>, float> pair_probs;
    for (const auto first : strings) {
      for (const auto second : strings) {
        if (first == second) {
          pair_probs[std::make_pair(first, second)] = 0;
        } else {
          float first_prob = probs[first];
          float second_prob = probs[second] / (1 - first_prob);
          pair_probs[std::make_pair(first, second)] = first_prob * second_prob;
        }
      }
    }

    std::map<std::string, float> inclusion_probs;
    for (const auto string : strings) {
      float inclusion_prob = 0.0;
      for (const auto other_string : strings) {
        inclusion_prob += pair_probs[std::make_pair(string, other_string)];
      }
      for (const auto other_string : strings) {
        inclusion_prob += pair_probs[std::make_pair(other_string, string)];
      }
      inclusion_probs[string] = inclusion_prob / 2;
    }

    int kTrials = 10000;

    std::vector<int> kNumSamples = {1, 2};

    for (const auto num_samples : kNumSamples) {
      std::map<std::string, int> counts;
      for (int i = 0; i < kTrials; i++) {
        auto nbests = lattice.NBest(num_samples, true, theta);
        for (const auto nbest : nbests) {
          counts[GetTokenized(nbest.first)]++;
        }
      }

      EXPECT_EQ(inclusion_probs.size(), counts.size());
      // If we take multiple samples WOR, we have to use corrected probs.
      std::map<std::string, float> probs_to_use =
          (num_samples == 1 ? probs : inclusion_probs);

      for (const auto &it : probs_to_use) {
        EXPECT_NEAR(it.second, 1.0 * counts[it.first] / (kTrials * num_samples),
                    0.02);
      }
    }
  }
}

TEST(LatticeTest, CalculateEntropyTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScore(&lattice, 0, 1, 0.0);  // A
  InsertWithScore(&lattice, 1, 1, 0.0);  // B
  InsertWithScore(&lattice, 2, 1, 0.1);  // C
  InsertWithScore(&lattice, 0, 2, 0.2);  // AB
  InsertWithScore(&lattice, 1, 2, 0.5);  // BC
  InsertWithScore(&lattice, 0, 3, 1.0);  // ABC

  // Calculate expected probabilities of each path
  const std::vector<double> kTheta = {0.0, 0.01, 0.5, 0.7, 1.0};
  for (const auto theta : kTheta) {
    std::vector<std::string> strings = {"ABC", "AB C", "A BC", "A B C"};
    std::map<std::string, float> probs;
    probs["ABC"] = std::exp(theta * 1.0);
    probs["AB C"] = std::exp(theta * (0.2 + 0.1));
    probs["A BC"] = std::exp(theta * (0.0 + 0.5));
    probs["A B C"] = std::exp(theta * (0.0 + 0.0 + 0.1));

    double Z = 0.0;
    for (const auto &it : probs) Z += it.second;
    for (auto &it : probs) it.second /= Z;

    for (const auto &it : strings) {
      EXPECT_EQ(1, probs.count(it));
    }
    float entropy = 0.0;
    for (const auto &it : probs) {
      entropy += (it.second * std::log(it.second));
    }
    EXPECT_NEAR(-entropy, lattice.CalculateEntropy(theta), 0.02);
  }
}

TEST(LatticeTest, ForwardAlgorithmTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScore(&lattice, 0, 1, 0.0);  // A
  InsertWithScore(&lattice, 1, 1, 0.0);  // B
  InsertWithScore(&lattice, 2, 1, 0.1);  // C
  InsertWithScore(&lattice, 0, 2, 0.2);  // AB
  InsertWithScore(&lattice, 1, 2, 0.5);  // BC
  InsertWithScore(&lattice, 0, 3, 1.0);  // ABC

  const std::vector<float> kTheta = {0.0, 0.01, 0.5, 0.7, 1.0};
  for (const auto theta : kTheta) {
    std::vector<float> alpha = lattice.ForwardAlgorithm(theta);
    EXPECT_EQ(alpha.size(), 8);  // 6 nodes, plus BOS, EOS
    // only alpha[C], alpha[EOS] have non-zero alpha
    for (int i : {0, 1, 2, 3}) {
      for (const auto &node : lattice.begin_nodes(i)) {
        if (i < 2) {
          EXPECT_EQ(alpha[node->node_id], 0.0);
        } else if (i == 2) {
          float Z =
              std::log(std::exp(theta * (0.0 + 0.0)) + std::exp(theta * 0.2));
          EXPECT_EQ(alpha[node->node_id], Z);
        } else if (i == 3) {
          float Z = std::log(std::exp(theta * (0.0 + 0.0 + 0.1)) +  // A + B + C
                             std::exp(theta * (0.2 + 0.1)) +        // AB + C
                             std::exp(theta * (0.0 + 0.5)) +        // A + BC
                             std::exp(theta * 1.0));                // ABC
          EXPECT_EQ(Z, alpha[node->node_id]);
        }
      }
    }
  }
}

TEST(LatticeTest, PopulateMarginalTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScoreAndId(&lattice, 0, 1, 1.0, 0);  // A
  InsertWithScoreAndId(&lattice, 1, 1, 1.2, 1);  // B
  InsertWithScoreAndId(&lattice, 2, 1, 2.5, 2);  // C
  InsertWithScoreAndId(&lattice, 0, 2, 3.0, 3);  // AB
  InsertWithScoreAndId(&lattice, 1, 2, 4.0, 4);  // BC
  InsertWithScoreAndId(&lattice, 0, 3, 2.0, 5);  // ABC

  std::vector<float> probs(6, 0.0);

  // Expand all paths:
  // A B C : exp(1.0 + 1.2 + 2.5) => path1
  // AB  C : exp(3.0 + 2.5)       => path2
  // A  BC : exp(1.0 + 4.0)       => path3
  // ABC   : exp(2.0)             => path4
  const float p1 = exp(1.0 + 1.2 + 2.5);
  const float p2 = exp(3.0 + 2.5);
  const float p3 = exp(1.0 + 4.0);
  const float p4 = exp(2.0);
  const float Z = p1 + p2 + p3 + p4;

  const float logZ = lattice.PopulateMarginal(1.0, &probs);

  EXPECT_NEAR((p1 + p3) / Z, probs[0], 0.001);  // A
  EXPECT_NEAR(p1 / Z, probs[1], 0.001);         // B
  EXPECT_NEAR((p1 + p2) / Z, probs[2], 0.001);  // C
  EXPECT_NEAR(p2 / Z, probs[3], 0.001);         // AB
  EXPECT_NEAR(p3 / Z, probs[4], 0.001);         // BC
  EXPECT_NEAR(p4 / Z, probs[5], 0.001);         // ABC
  EXPECT_NEAR(std::log(static_cast<double>(Z)), logZ, 0.001);
}

TEST(LatticeTest, SampleTest) {
  Lattice lattice;
  lattice.SetSentence("ABC");

  InsertWithScoreAndId(&lattice, 0, 1, 1.0, 0);  // A
  InsertWithScoreAndId(&lattice, 1, 1, 1.2, 1);  // B
  InsertWithScoreAndId(&lattice, 2, 1, 1.5, 2);  // C
  InsertWithScoreAndId(&lattice, 0, 2, 1.6, 3);  // AB
  InsertWithScoreAndId(&lattice, 1, 2, 1.7, 4);  // BC
  InsertWithScoreAndId(&lattice, 0, 3, 1.8, 5);  // ABC

  const std::vector<double> kTheta = {0.0, 0.01, 0.5, 0.7, 1.0};
  for (int i = 0; i < kTheta.size(); ++i) {
    std::map<std::string, double> probs;
    // Expands all paths in the lattice.
    probs["A B C"] = exp(kTheta[i] * (1.0 + 1.2 + 1.5));  // A B C
    probs["AB C"] = exp(kTheta[i] * (1.6 + 1.5));         // AB C
    probs["A BC"] = exp(kTheta[i] * (1.0 + 1.7));         // A BC
    probs["ABC"] = exp(kTheta[i] * 1.8);                  // ABC

    // Computes expected probabilities.
    double Z = 0.0;
    for (const auto &it : probs) Z += it.second;
    for (auto &it : probs) it.second /= Z;

    // Samples `kTrial` times and verifies the probabilities.
    constexpr int kTrial = 100000;
    std::map<std::string, int> freq;
    for (int n = 0; n < kTrial; ++n) {
      freq[GetTokenized(lattice.Sample(kTheta[i]))]++;
    }

    EXPECT_EQ(probs.size(), freq.size());
    for (const auto &it : probs) {
      EXPECT_NEAR(it.second, 1.0 * freq[it.first] / kTrial, 0.02);
    }
  }
}

ModelProto MakeBaseModelProto() {
  ModelProto model_proto;
  auto *sp1 = model_proto.add_pieces();
  auto *sp2 = model_proto.add_pieces();
  auto *sp3 = model_proto.add_pieces();

  sp1->set_type(ModelProto::SentencePiece::UNKNOWN);
  sp1->set_piece("<unk>");
  sp2->set_type(ModelProto::SentencePiece::CONTROL);
  sp2->set_piece("<s>");
  sp3->set_type(ModelProto::SentencePiece::CONTROL);
  sp3->set_piece("</s>");

  return model_proto;
}

// Returns model protos in parameterized tests.
const std::vector<EncoderVersion> &GetEncoderVersions() {
  static const std::vector<EncoderVersion> &v =
      *new std::vector<EncoderVersion>{EncoderVersion::kOptimized,
                                       EncoderVersion::kOriginal};
  return v;
}

class UnigramModelTest : public test::TestWithParam<EncoderVersion> {
 protected:
  void SetUp() override { encoder_version_ = GetParam(); }
  void TearDown() override {}
  EncoderVersion encoder_version_;
};

void AddPiece(ModelProto *model_proto, const std::string &piece,
              float score = 0.0) {
  auto *sp = model_proto->add_pieces();
  sp->set_piece(piece);
  sp->set_score(score);
}

TEST(UnigramModelTest, SetUnigramModelTest) {
  ModelProto model_proto = MakeBaseModelProto();

  AddPiece(&model_proto, "a");
  AddPiece(&model_proto, "b");
  AddPiece(&model_proto, "c");
  AddPiece(&model_proto, "d");

  const Model model(model_proto);
  EXPECT_EQ(model_proto.SerializeAsString(),
            model.model_proto().SerializeAsString());
}

TEST(UnigramModelTest, SampleEncodeAndScoreTest) {
  // Test whether inclusion probabilities are correct
  ModelProto model_proto = MakeBaseModelProto();
  AddPiece(&model_proto, "A", 0.0);    // 3
  AddPiece(&model_proto, "B", 0.0);    // 4
  AddPiece(&model_proto, "C", 0.1);    // 5
  AddPiece(&model_proto, "AB", 0.2);   // 6
  AddPiece(&model_proto, "BC", 0.5);   // 7
  AddPiece(&model_proto, "ABC", 1.0);  // 8

  Model model(model_proto);

  Lattice lattice;
  lattice.SetSentence("ABC");
  model.PopulateNodes(&lattice);

  std::vector<float> kTheta = {0.0, 1.0};

  for (const auto theta : kTheta) {
    std::vector<std::string> strings = {"ABC", "AB C", "A BC", "A B C"};
    std::map<std::string, float> probs;
    probs["ABC"] = std::exp(theta * 1.0);
    probs["AB C"] = std::exp(theta * (0.2 + 0.1));
    probs["A BC"] = std::exp(theta * (0.0 + 0.5));
    probs["A B C"] = std::exp(theta * (0.0 + 0.0 + 0.1));

    for (const auto &it : strings) {
      EXPECT_EQ(1, probs.count(it));
    }

    double Z = 0.0;
    for (const auto &it : probs) Z += it.second;
    for (auto &it : probs) it.second /= Z;

    std::map<std::pair<std::string, std::string>, float> pair_probs;
    for (const auto first : strings) {
      for (const auto second : strings) {
        if (first == second) {
          pair_probs[std::make_pair(first, second)] = 0;
        } else {
          float first_prob = probs[first];
          float second_prob = probs[second] / (1 - first_prob);
          pair_probs[std::make_pair(first, second)] = first_prob * second_prob;
        }
      }
    }

    std::map<std::string, float> inclusion_probs;
    for (const auto string : strings) {
      float inclusion_prob = 0.0;
      for (const auto other_string : strings) {
        inclusion_prob += pair_probs[std::make_pair(string, other_string)];
      }
      for (const auto other_string : strings) {
        inclusion_prob += pair_probs[std::make_pair(other_string, string)];
      }
      inclusion_probs[string] = inclusion_prob / 2;
    }
    std::vector<int> kNumSamples = {1, 2};

    for (const auto num_samples : kNumSamples) {
      std::map<std::string, int> counts;
      std::map<std::string, float> scores;
      int kTrials = 50000;
      for (int i = 0; i < kTrials; i++) {
        NBestEncodeResult sample =
            model.SampleEncodeAndScore("ABC", theta, num_samples, true, false);

        for (const auto &it : sample) {
          std::vector<std::string> tokens;
          for (const auto &inner_it : it.first) {
            tokens.push_back(std::string(inner_it.first));
          }
          std::string sample_string = absl::StrJoin(tokens, " ");
          counts[sample_string] += 1;
          // use the fact that E(1_{i in sample} / score of i) = 1
          // see https://arxiv.org/pdf/1903.06059.pdf appendix D
          scores[sample_string] += std::exp(-it.second);
        }
      }

      // Check that counts and probs are correct
      std::map<std::string, float> probs_to_use =
          (num_samples == 1 ? probs : inclusion_probs);

      for (const auto &it : scores) Z += it.second;
      for (const auto &it : probs_to_use) {
        EXPECT_NEAR(it.second, 1.0 * counts[it.first] / (kTrials * num_samples),
                    0.02);
        // The expectation is quite loose, use a higher tolerance
        EXPECT_NEAR(1.0, scores[it.first] / kTrials, 0.20);
      }
    }
  }
}

TEST_P(UnigramModelTest, PieceToIdTest) {
  ModelProto model_proto = MakeBaseModelProto();

  AddPiece(&model_proto, "a", 0.1);
  AddPiece(&model_proto, "b", 0.2);
  AddPiece(&model_proto, "c", 0.3);
  AddPiece(&model_proto, "d", 0.4);

  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());

  EXPECT_EQ(model_proto.SerializeAsString(),
            model.model_proto().SerializeAsString());

  EXPECT_NEAR(0.1, model.min_score(), 0.001);
  EXPECT_NEAR(0.4, model.max_score(), 0.001);

  EXPECT_EQ(0, model.PieceToId("<unk>"));
  EXPECT_EQ(1, model.PieceToId("<s>"));
  EXPECT_EQ(2, model.PieceToId("</s>"));
  EXPECT_EQ(3, model.PieceToId("a"));
  EXPECT_EQ(3, model.PieceToId(absl::string_view("a b", 1)));
  EXPECT_EQ(4, model.PieceToId("b"));
  EXPECT_EQ(5, model.PieceToId("c"));
  EXPECT_EQ(6, model.PieceToId("d"));
  EXPECT_EQ(0, model.PieceToId("e"));  // unk
  EXPECT_EQ(0, model.PieceToId(""));   // unk

  EXPECT_EQ("<unk>", model.IdToPiece(0));
  EXPECT_EQ("<s>", model.IdToPiece(1));
  EXPECT_EQ("</s>", model.IdToPiece(2));
  EXPECT_EQ("a", model.IdToPiece(3));
  EXPECT_EQ("b", model.IdToPiece(4));
  EXPECT_EQ("c", model.IdToPiece(5));
  EXPECT_EQ("d", model.IdToPiece(6));

  EXPECT_TRUE(model.IsUnknown(0));
  EXPECT_FALSE(model.IsUnknown(1));
  EXPECT_FALSE(model.IsUnknown(2));
  EXPECT_FALSE(model.IsUnknown(3));
  EXPECT_FALSE(model.IsUnknown(4));
  EXPECT_FALSE(model.IsUnknown(5));
  EXPECT_FALSE(model.IsUnknown(6));

  EXPECT_FALSE(model.IsControl(0));
  EXPECT_TRUE(model.IsControl(1));
  EXPECT_TRUE(model.IsControl(2));
  EXPECT_FALSE(model.IsControl(3));
  EXPECT_FALSE(model.IsControl(4));
  EXPECT_FALSE(model.IsControl(5));
  EXPECT_FALSE(model.IsControl(6));

  EXPECT_NEAR(0, model.GetScore(0), 0.0001);
  EXPECT_NEAR(0, model.GetScore(1), 0.0001);
  EXPECT_NEAR(0, model.GetScore(2), 0.0001);
  EXPECT_NEAR(0.1, model.GetScore(3), 0.0001);
  EXPECT_NEAR(0.2, model.GetScore(4), 0.0001);
  EXPECT_NEAR(0.3, model.GetScore(5), 0.0001);
  EXPECT_NEAR(0.4, model.GetScore(6), 0.0001);

  EXPECT_TRUE(model.Encode("").empty());
}

TEST_P(UnigramModelTest, PopulateNodesAllUnknownsTest) {
  ModelProto model_proto = MakeBaseModelProto();
  AddPiece(&model_proto, "x");
  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());

  Lattice lattice;
  lattice.SetSentence("abc");
  model.PopulateNodes(&lattice);

  EXPECT_EQ(1, lattice.begin_nodes(0).size());
  EXPECT_EQ(1, lattice.begin_nodes(1).size());
  EXPECT_EQ(1, lattice.begin_nodes(2).size());

  EXPECT_EQ(0, lattice.begin_nodes(0)[0]->id);
  EXPECT_EQ(0, lattice.begin_nodes(1)[0]->id);
  EXPECT_EQ(0, lattice.begin_nodes(2)[0]->id);
}

TEST_P(UnigramModelTest, PopulateNodesTest) {
  ModelProto model_proto = MakeBaseModelProto();

  AddPiece(&model_proto, "a", 0.1);   // 3
  AddPiece(&model_proto, "b", 0.2);   // 4
  AddPiece(&model_proto, "ab", 0.3);  // 5
  AddPiece(&model_proto, "bc", 0.4);  // 6

  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());

  Lattice lattice;
  lattice.SetSentence("abc");

  model.PopulateNodes(&lattice);

  EXPECT_EQ(2, lattice.begin_nodes(0).size());  // a,ab
  EXPECT_EQ(2, lattice.begin_nodes(1).size());  // b,bc
  EXPECT_EQ(1, lattice.begin_nodes(2).size());  // c(unk)

  EXPECT_EQ(3, lattice.begin_nodes(0)[0]->id);
  EXPECT_EQ(5, lattice.begin_nodes(0)[1]->id);
  EXPECT_EQ(4, lattice.begin_nodes(1)[0]->id);
  EXPECT_EQ(6, lattice.begin_nodes(1)[1]->id);
  EXPECT_EQ(0, lattice.begin_nodes(2)[0]->id);

  EXPECT_NEAR(0.1, lattice.begin_nodes(0)[0]->score, 0.001);
  EXPECT_NEAR(0.3, lattice.begin_nodes(0)[1]->score, 0.001);
  EXPECT_NEAR(0.2, lattice.begin_nodes(1)[0]->score, 0.001);
  EXPECT_NEAR(0.4, lattice.begin_nodes(1)[1]->score, 0.001);
}

TEST_P(UnigramModelTest, PopulateNodesWithUnusedTest) {
  ModelProto model_proto = MakeBaseModelProto();

  AddPiece(&model_proto, "a", 0.1);   // 3
  AddPiece(&model_proto, "b", 0.2);   // 4
  AddPiece(&model_proto, "ab", 0.3);  // 5
  AddPiece(&model_proto, "bc", 0.4);  // 6

  model_proto.mutable_pieces(5)->set_type(ModelProto::SentencePiece::UNUSED);
  model_proto.mutable_pieces(6)->set_type(ModelProto::SentencePiece::UNUSED);

  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());

  Lattice lattice;
  lattice.SetSentence("abc");

  model.PopulateNodes(&lattice);

  EXPECT_EQ(1, lattice.begin_nodes(0).size());  // a
  EXPECT_EQ(1, lattice.begin_nodes(1).size());  // b
  EXPECT_EQ(1, lattice.begin_nodes(2).size());  // c(unk)
  EXPECT_EQ(3, lattice.begin_nodes(0)[0]->id);
  EXPECT_EQ(4, lattice.begin_nodes(1)[0]->id);
  EXPECT_EQ(0, lattice.begin_nodes(2)[0]->id);
}

TEST_P(UnigramModelTest, ModelNBestTest) {
  ModelProto model_proto = MakeBaseModelProto();
  AddPiece(&model_proto, "a", 0.0);     // 3
  AddPiece(&model_proto, "b", 0.0);     // 4
  AddPiece(&model_proto, "c", 0.0);     // 5
  AddPiece(&model_proto, "ab", 2.0);    // 6
  AddPiece(&model_proto, "bc", 5.0);    // 7
  AddPiece(&model_proto, "abc", 10.0);  // 8

  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());

  auto nbest = model.NBestEncode("", 10);
  EXPECT_EQ(1, nbest.size());
  EXPECT_TRUE(nbest[0].first.empty());
  EXPECT_EQ(0.0, nbest[0].second);

  nbest = model.NBestEncode("abc", 10);
  EXPECT_EQ(4, nbest.size());

  auto sample = model.SampleEncode("", 0.1);
  EXPECT_EQ(0, sample.size());
  sample = model.SampleEncode("abc", 0.1);
  EXPECT_FALSE(sample.empty());
}

TEST_P(UnigramModelTest, EncodeTest) {
  ModelProto model_proto = MakeBaseModelProto();
  AddPiece(&model_proto, "ab", 0.0);         // 3
  AddPiece(&model_proto, "cd", -0.1);        // 4
  AddPiece(&model_proto, "abc", -0.2);       // 5
  AddPiece(&model_proto, "a", -0.3);         // 6
  AddPiece(&model_proto, "b", -0.4);         // 7
  AddPiece(&model_proto, "c", -0.5);         // 8
  AddPiece(&model_proto, "ABC", -0.5);       // 9
  AddPiece(&model_proto, "abcdabcd", -0.5);  // 10
  AddPiece(&model_proto, "q", -0.5);         // 11
  AddPiece(&model_proto, "r", -0.5);         // 12
  AddPiece(&model_proto, "qr", -0.5);        // 13
  model_proto.mutable_pieces(9)->set_type(   // ABC
      ModelProto::SentencePiece::USER_DEFINED);
  model_proto.mutable_pieces(10)->set_type(  // abcdabcd
      ModelProto::SentencePiece::USER_DEFINED);
  model_proto.mutable_pieces(11)->set_type(  // q
      ModelProto::SentencePiece::USER_DEFINED);
  model_proto.mutable_pieces(12)->set_type(  // r
      ModelProto::SentencePiece::USER_DEFINED);

  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());

  EncodeResult result;

  result = model.Encode("abc");
  EXPECT_EQ(1, result.size());
  EXPECT_EQ("abc", result[0].first);

  result = model.Encode("AB");
  EXPECT_EQ(2, result.size());
  EXPECT_EQ("A", result[0].first);
  EXPECT_EQ("B", result[1].first);

  result = model.Encode("abcd");
  EXPECT_EQ(2, result.size());
  EXPECT_EQ("ab", result[0].first);
  EXPECT_EQ("cd", result[1].first);

  result = model.Encode("abcc");
  EXPECT_EQ(2, result.size());
  EXPECT_EQ("abc", result[0].first);
  EXPECT_EQ("c", result[1].first);

  result = model.Encode("xabcabaabcdd");
  EXPECT_EQ(7, result.size());
  EXPECT_EQ("x", result[0].first);
  EXPECT_EQ("abc", result[1].first);
  EXPECT_EQ("ab", result[2].first);
  EXPECT_EQ("a", result[3].first);
  EXPECT_EQ("ab", result[4].first);
  EXPECT_EQ("cd", result[5].first);
  EXPECT_EQ("d", result[6].first);

  // all unknown.
  result = model.Encode("xyz東京");
  EXPECT_EQ(5, result.size());
  EXPECT_EQ("x", result[0].first);
  EXPECT_EQ("y", result[1].first);
  EXPECT_EQ("z", result[2].first);
  EXPECT_EQ("東", result[3].first);
  EXPECT_EQ("京", result[4].first);

  // User defined
  result = model.Encode("ABC");
  EXPECT_EQ(1, result.size());
  EXPECT_EQ("ABC", result[0].first);

  result = model.Encode("abABCcd");
  EXPECT_EQ(3, result.size());
  EXPECT_EQ("ab", result[0].first);
  EXPECT_EQ("ABC", result[1].first);
  EXPECT_EQ("cd", result[2].first);

  // middle "abcdabcd" is user defined.
  result = model.Encode("ababcdabcdcd");
  EXPECT_EQ(3, result.size());
  EXPECT_EQ("ab", result[0].first);
  EXPECT_EQ("abcdabcd", result[1].first);
  EXPECT_EQ("cd", result[2].first);

  result = model.Encode("abqrcd");
  EXPECT_EQ(4, result.size());
  EXPECT_EQ("ab", result[0].first);
  EXPECT_EQ("q", result[1].first);
  EXPECT_EQ("r", result[2].first);
  EXPECT_EQ("cd", result[3].first);
}

TEST_P(UnigramModelTest, EncodeWithUnusedTest) {
  ModelProto model_proto = MakeBaseModelProto();

  AddPiece(&model_proto, "abcd", 10.0);  // 3
  AddPiece(&model_proto, "abc", 5.0);    // 4
  AddPiece(&model_proto, "ab", 2.0);     // 5
  AddPiece(&model_proto, "cd", 1.0);     // 6
  AddPiece(&model_proto, "a", 0.0);      // 7
  AddPiece(&model_proto, "b", 0.0);      // 8
  AddPiece(&model_proto, "c", 0.0);      // 9
  AddPiece(&model_proto, "d", 0.0);      // 10

  // No unused.
  {
    Model model(model_proto);
    EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());
    const auto result = model.Encode("abcd");
    EXPECT_EQ(1, result.size());
    EXPECT_EQ("abcd", result[0].first);
  }

  {
    model_proto.mutable_pieces(3)->set_type(ModelProto::SentencePiece::UNUSED);
    Model model(model_proto);
    EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());
    const auto result = model.Encode("abcd");
    EXPECT_EQ(2, result.size());
    EXPECT_EQ("abc", result[0].first);
    EXPECT_EQ("d", result[1].first);
  }

  {
    model_proto.mutable_pieces(3)->set_type(ModelProto::SentencePiece::UNUSED);
    model_proto.mutable_pieces(5)->set_type(ModelProto::SentencePiece::UNUSED);
    Model model(model_proto);
    EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());
    const auto result = model.Encode("abcd");
    EXPECT_EQ(2, result.size());
    EXPECT_EQ("abc", result[0].first);
    EXPECT_EQ("d", result[1].first);
  }

  {
    // This is different from BPE segmentation.
    // Unigram language model simply finds the best path without unused nodes.
    model_proto.mutable_pieces(3)->set_type(ModelProto::SentencePiece::UNUSED);
    model_proto.mutable_pieces(4)->set_type(ModelProto::SentencePiece::UNUSED);
    model_proto.mutable_pieces(5)->set_type(ModelProto::SentencePiece::NORMAL);
    Model model(model_proto);
    EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());
    const auto result = model.Encode("abcd");
    EXPECT_EQ(2, result.size());
    EXPECT_EQ("ab", result[0].first);
    EXPECT_EQ("cd", result[1].first);
  }
}

TEST_P(UnigramModelTest, VerifyOutputsEquivalent) {
  ModelProto model_proto = MakeBaseModelProto();

  AddPiece(&model_proto, "abcd", 10.0);  // 3
  AddPiece(&model_proto, "abc", 5.0);    // 4
  AddPiece(&model_proto, "ab", 6.0);     // 5
  AddPiece(&model_proto, "cd", 4.0);     // 6
  AddPiece(&model_proto, "a", 4.0);      // 7
  AddPiece(&model_proto, "b", 1.9);      // 8
  AddPiece(&model_proto, "c", 2.0);      // 9
  AddPiece(&model_proto, "d", 1.0);      // 10
  Model model(model_proto);
  EXPECT_TRUE(model.SetEncoderVersion(encoder_version_).ok());
  // Equivalent outputs.
  EXPECT_TRUE(model.VerifyOutputsEquivalent("", ""));
  EXPECT_TRUE(model.VerifyOutputsEquivalent("a b", "a b"));
  EXPECT_TRUE(model.VerifyOutputsEquivalent("abcd", "ab cd"));

  // Inequivalent outputs.
  EXPECT_FALSE(model.VerifyOutputsEquivalent("a", "a b"));
  EXPECT_FALSE(model.VerifyOutputsEquivalent("ab", "a b"));
}

INSTANTIATE_TEST_SUITE_P(ParametrizedUnigramModelTests, UnigramModelTest,
                         test::ValuesIn(GetEncoderVersions()));

}  // namespace unigram
}  // namespace sentencepiece
