// 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.!

#ifndef UNIGRAM_MODEL_H_
#define UNIGRAM_MODEL_H_

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "common.h"
#include "freelist.h"
#include "model_interface.h"
#include "sentencepiece_model.pb.h"
#include "third_party/darts_clone/darts.h"

namespace sentencepiece {
namespace unigram {

// Lattice represents a search space of sentence piece segmentation.
class Lattice {
 public:
  Lattice();
  virtual ~Lattice();

  struct Node {
    absl::string_view piece;  // Sentence piece representation.
    uint32_t pos;               // Unicode position in the sentence.
    uint32_t length;            // Unicode length, not UT8 byte.
    uint32_t node_id;           // unique id in the current lattice.
    int id;                   // vocab id. (maybe -1 for UNK)
    float score;              // logprob of this sentencepiece.
    float backtrace_score;    // backtrace info used in Viterbi.
    Node *prev;               // best previous node on Viterbi path.

    std::string DebugString() const;
  };

  // Returns bos node.
  Node *bos_node() const;

  // Returns eos node.
  Node *eos_node() const;

  // Returns nodes starting at |pos|.
  const std::vector<Node *> &begin_nodes(int pos) const;

  // Returns nodes ending at |pos|.
  const std::vector<Node *> &end_nodes(int pos) const;

  // Returns Unicode character length.
  int size() const;

  // Returns multi-byte (utf8) length.
  int utf8_size() const;

  // Returns the substring of sentence. sentence[pos:]
  const char *surface(int pos) const;

  // Returns immutable sentence. The same as surface(0)
  const char *sentence() const;

  // Clears the lattice.
  void Clear();

  // Sets new sentence.
  void SetSentence(absl::string_view sentence);

  // Inserts a new node at [pos, pos + length - 1].
  // After calling this method, The caller must set Node::score and Node::id.
  Node *Insert(int pos, int length);

  using LatticePathWithScore = std::pair<std::vector<Node *>, float>;

  // Returns Viterbi path. All nodes must be populated in advance.
  LatticePathWithScore Viterbi();

  // Runs forwards/backwards algorithm, returns vector with normalised
  // transition probs.
  std::vector<float> ForwardAlgorithm(float theta) const;
  std::vector<float> BackwardAlgorithm(float theta) const;

  // Returns n-best results.
  std::vector<LatticePathWithScore> NBest(size_t nbest_size, bool sample,
                                          float theta);

  // Samples one path from the lattice according to the
  // generation probability (Product of piece probabilities).
  // `theta` is a smoothing parameter.
  std::vector<Node *> Sample(float theta);

  // Calculates the entropy of the lattice.
  float CalculateEntropy(float theta) const;

  // Populates marginal probability of every node in this lattice.
  // |freq| is the frequency of the sentence.
  //  for (auto *node : all_nodes_) {
  //    (*expected)[node->id] += marginal_prob_of_node * freq;
  //  }
  // Returns the log-likelihood of this sentence.
  float PopulateMarginal(float freq, std::vector<float> *expected) const;

 private:
  // Returns new node.
  // Lattice class has the ownership of the returned value.
  Node *NewNode();

  absl::string_view sentence_;
  std::vector<const char *> surface_;
  std::vector<std::vector<Node *>> begin_nodes_;
  std::vector<std::vector<Node *>> end_nodes_;
  model::FreeList<Node> node_allocator_;
};

class Model : public ModelInterface {
 public:
  explicit Model(const ModelProto &model_proto);
  Model() {}
  ~Model() override;

  EncodeResult Encode(absl::string_view normalized) const override;

  NBestEncodeResult NBestEncode(absl::string_view normalized,
                                int nbest_size) const override;

  EncodeResult SampleEncode(absl::string_view normalized,
                            float theta) const override;

  NBestEncodeResult SampleEncodeAndScore(absl::string_view normalized,
                                         float theta, int samples, bool wor,
                                         bool include_best) const override;

  float CalculateEntropy(absl::string_view normalized,
                         float theta) const override;

  bool IsSampleEncodeAvailable() const override { return true; }

  bool IsSampleEncodeAndScoreAvailable() const override { return true; }

  bool IsCalculateEntropyAvailable() const override { return true; }

  bool IsNBestEncodeAvailable() const override { return true; }

  // Returns the minimum score in sentence pieces.
  // min_score() - 10 is used for the cost of unknown sentence.
  float min_score() const { return min_score_; }

  // Returns the maximum score in sentence pieces.
  // max_score() is used for the cost of user defined symbols.
  float max_score() const { return max_score_; }

  // Populates all sentence pieces to the |lattice|.
  // After calling this function, lattice.Viterbi() returns the
  // best segmentation.
  void PopulateNodes(Lattice *lattice) const;

  // Returns a vocab id of |piece|.
  int PieceToId(absl::string_view piece) const override;

  // Verifies if two outputs are equivalent by comparing their scores.
  bool VerifyOutputsEquivalent(absl::string_view expected,
                               absl::string_view actual) const override;

  enum EncoderVersion {
    kOptimized,  // The optimized encoder.
    kOriginal    // The original encoder.
  };

  void SetEncoderVersion(EncoderVersion encoder_version) {
    encoder_version_ = encoder_version;
  }

  // Returns the current encoder version in use.
  EncoderVersion GetEncoderVersion() const { return encoder_version_; }

 protected:
  // Builds a Trie index.
  void BuildTrie(std::vector<std::pair<absl::string_view, int>> *pieces);

  // The optimized Viterbi encode.
  // Main differences from the original function:
  // 1. Memorizes the best path at each postion so far,
  // 2. No need to store the Lattice nodes,
  // 3. Works in utf-8 directly,
  // 4. Defines a new struct with fewer fields than Lattice,
  // 5. Does not depend on `class Lattice` nor call `SetSentence()`,
  // `PopulateNodes()`, or `Viterbi()`. It does everything in one function.
  // For detailed explanations please see the comments inside the function body.
  EncodeResult EncodeOptimized(absl::string_view normalized) const;

  float min_score_ = 0.0;
  float max_score_ = 0.0;
  std::unique_ptr<Darts::DoubleArray> trie_;

  // Maximum size of the return value of Trie, which corresponds
  // to the maximum size of shared common prefix in the sentence pieces.
  int trie_results_size_;

  // encoder version.
  EncoderVersion encoder_version_ = kOptimized;
};

}  // namespace unigram
}  // namespace sentencepiece
#endif  // UNIGRAM_MODEL_H_